This repository has been archived by the owner on Oct 12, 2018. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathrangeloss.py
115 lines (85 loc) · 3.78 KB
/
rangeloss.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import mxnet as mx
import numpy as np
class RangeLossOp(mx.operator.CustomOp):
def __init__(self, margin, k, alpha, beta):
self.margin = margin
self.k = k
self.alpha = alpha
self.beta = beta
def compute_top_k(self, features):
num = features.shape[0]
dists = np.zeros((num, num))
for id_1 in range(0, num):
for id_2 in range(id_1+1, num):
dist = np.linalg.norm(features[id_1] - features[id_2])
dists[id_1, id_2] = dist
dist_array = dists.reshape((1, -1))
return dist_array.sort()[-self.k:]
def compute_min_dist(self, centers):
num = centers.shape[0]
dists = np.ones((num, num)) * -1
for id_1 in range(0, num):
for id_2 in range(id_1+1, num):
dist = np.linalg.norm(centers[id_1] - centers[id_2])
dists[id_1, id_2] = dist
dist_array = dists.reshape((1, -1))
dist_array = dist_array[np.where(dist_array > 0)]
return dist_array.sort()[0]
def forward(self, is_train, req, in_data, out_data, aux):
features = in_data[0].asnumpy()
labels = in_data[1].asnumpy()
unique_labels, counts = np.unique(labels, return_counts=True)
centers = np.zeros((unique_labels.shape[0], features.shape[1]))
d = np.zeros((unique_labels.shape[0], self.k))
l_r = np.zeros((unique_labels.shape[0]))
for idx, l in enumerate(unique_labels):
indices = np.where(labels == l)
features_l = features[indices, :]
centers[idx, :] = np.mean(features_l, axis=0)
d[idx, :] = self.compute_top_k(features_l)
l_r[idx] = self.k / np.sum(d[idx, :])
l_intra = np.sum(l_r)
d_center = self.compute_min_dist(centers)
l_inter = max(self.margin - d_center, 0)
loss = l_intra * self.alpha + l_inter * self.beta
self.assign(out_data[0], req[0], mx.nd.array(loss))
def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
features = in_data[0].asnumpy()
labels = in_data[1].asnumpy().ravel().astype(np.int)
unique_labels, counts = np.unique(labels, return_counts=True)
centers = np.zeros((unique_labels.shape[0], features.shape[1]))
d = np.zeros((unique_labels.shape[0], self.k))
l_r = np.zeros((unique_labels.shape[0]))
for idx, l in enumerate(unique_labels):
indices = np.where(labels == l)
features_l = features[indices, :]
centers[idx, :] = np.mean(features_l, axis=0)
d[idx, :] = self.compute_top_k(features_l)
l_r[idx] = self.k / np.sum(d[idx, :])
l_intra = np.sum(l_r)
d_center = self.compute_min_dist(centers)
l_inter = max(self.margin - d_center, 0)
self.assign(in_grad[0], req[0], mx.nd.array(y))
@mx.operator.register('rangeloss')
class RangeLossOpProp(mx.operator.CustomOpProp):
def __init__(self, num_hidden, alpha, beta, margin=20000, k=2):
super(RangeLossOpProp, self).__init__(need_top_grad=False)
self.margin = margin
self.num_hidden = int(num_hidden)
self.alpha = float(alpha)
self.beta = float(beta)
self.k = int(k)
def list_arguments(self):
return ['data', 'label']
def list_outputs(self):
return ['output']
def infer_shape(self, in_shape):
data_shape = in_shape[0]
label_shape = (in_shape[0][0], )
out_shape = in_shape[0]
return [data_shape, label_shape], [out_shape], []
def infer_type(self, in_type):
dtype = in_type[0]
return [dtype, dtype], [dtype], []
def create_operator(self, ctx, shapes, dtypes):
return RangeLossOp(margin=self.margin, alpha=self.alpha, beta=self.beta, k=self.k)