-
Notifications
You must be signed in to change notification settings - Fork 2
/
ArcFace_losses.py
127 lines (109 loc) · 5.75 KB
/
ArcFace_losses.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
116
117
118
119
120
121
122
123
124
125
126
127
import tensorflow as tf
import math
def arcface_loss(embedding, labels, out_num, w_init=None, s=64., m=0.5):
'''
:param embedding: the input embedding vectors
:param labels: the input labels, the shape should be eg: (batch_size, 1)
:param s: scalar value default is 64
:param out_num: output class num
:param m: the margin value, default is 0.5
:return: the final cacualted output, this output is send into the tf.nn.softmax directly
'''
cos_m = math.cos(m)
sin_m = math.sin(m)
mm = sin_m * m # issue 1
threshold = math.cos(math.pi - m)
with tf.variable_scope('arcface_loss'):
# inputs and weights norm
ori_norm_embd = tf.norm(embedding, axis=1, keep_dims=True)
embedding_norm = tf.maximum(ori_norm_embd, 1e-6)
embedding = tf.div(embedding, embedding_norm, name='norm_embedding')
weights = tf.get_variable(name='embedding_weights', shape=(embedding.get_shape().as_list()[-1], out_num),
initializer=w_init, dtype=tf.float32)
ori_norm_w = tf.norm(weights, axis=0, keep_dims=True)
weights_norm = tf.maximum(ori_norm_w, 1e-6)
weights = tf.div(weights, weights_norm, name='norm_weights')
# cos(theta+m)
cos_t = tf.matmul(embedding, weights, name='cos_t')
cos_t2 = tf.square(cos_t, name='cos_2')
sin_t2 = tf.subtract(1., cos_t2, name='sin_2')
sin_t = tf.sqrt(sin_t2, name='sin_t')
cos_mt = s * tf.subtract(tf.multiply(cos_t, cos_m), tf.multiply(sin_t, sin_m), name='cos_mt')
# this condition controls the theta+m should in range [0, pi]
# 0<=theta+m<=pi
# -m<=theta<=pi-m
cond_v = cos_t - threshold
cond = tf.cast(tf.nn.relu(cond_v, name='if_else'), dtype=tf.bool)
keep_val = s*(cos_t - mm)
cos_mt_temp = tf.where(cond, cos_mt, keep_val)
mask = tf.one_hot(labels, depth=out_num, name='one_hot_mask')
# mask = tf.squeeze(mask, 1)
inv_mask = tf.subtract(1., mask, name='inverse_mask')
s_cos_t = tf.multiply(s, cos_t, name='scalar_cos_t')
output = tf.add(tf.multiply(s_cos_t, inv_mask), tf.multiply(cos_mt_temp, mask), name='arcface_loss_output')
return output, ori_norm_embd, ori_norm_w
def cosineface_losses(embedding, labels, out_num, w_init=None, s=30., m=0.4):
'''
:param embedding: the input embedding vectors
:param labels: the input labels, the shape should be eg: (batch_size, 1)
:param s: scalar value, default is 30
:param out_num: output class num
:param m: the margin value, default is 0.4
:return: the final cacualted output, this output is send into the tf.nn.softmax directly
'''
with tf.variable_scope('cosineface_loss'):
# inputs and weights norm
embedding_norm = tf.norm(embedding, axis=1, keep_dims=True)
embedding = tf.div(embedding, embedding_norm, name='norm_embedding')
weights = tf.get_variable(name='embedding_weights', shape=(embedding.get_shape().as_list()[-1], out_num),
initializer=w_init, dtype=tf.float32)
weights_norm = tf.norm(weights, axis=0, keep_dims=True)
weights = tf.div(weights, weights_norm, name='norm_weights')
# cos_theta - m
cos_t = tf.matmul(embedding, weights, name='cos_t')
cos_t_m = tf.subtract(cos_t, m, name='cos_t_m')
mask = tf.one_hot(labels, depth=out_num, name='one_hot_mask')
inv_mask = tf.subtract(1., mask, name='inverse_mask')
output = tf.add(s * tf.multiply(cos_t, inv_mask), s * tf.multiply(cos_t_m, mask), name='cosineface_loss_output')
return output
def combine_loss_val(embedding, labels, w_init, out_num, margin_a, margin_m, margin_b, s):
'''
This code is contributed by RogerLo. Thanks for you contribution.
:param embedding: the input embedding vectors
:param labels: the input labels, the shape should be eg: (batch_size, 1)
:param s: scalar value default is 64
:param out_num: output class num
:param m: the margin value, default is 0.5
:return: the final cacualted output, this output is send into the tf.nn.softmax directly
'''
weights = tf.get_variable(name='embedding_weights', shape=(embedding.get_shape().as_list()[-1], out_num),
initializer=w_init, dtype=tf.float32)
weights_unit = tf.nn.l2_normalize(weights, axis=0)
embedding_unit = tf.nn.l2_normalize(embedding, axis=1)
cos_t = tf.matmul(embedding_unit, weights_unit)
ordinal = tf.constant(list(range(0, embedding.get_shape().as_list()[0])), tf.int64)
ordinal_y = tf.stack([ordinal, labels], axis=1)
zy = cos_t * s
sel_cos_t = tf.gather_nd(zy, ordinal_y)
if margin_a != 1.0 or margin_m != 0.0 or margin_b != 0.0:
if margin_a == 1.0 and margin_m == 0.0:
s_m = s * margin_b
new_zy = sel_cos_t - s_m
else:
cos_value = sel_cos_t / s
t = tf.acos(cos_value)
if margin_a != 1.0:
t = t * margin_a
if margin_m > 0.0:
t = t + margin_m
body = tf.cos(t)
if margin_b > 0.0:
body = body - margin_b
new_zy = body * s
updated_logits = tf.add(zy, tf.scatter_nd(ordinal_y, tf.subtract(new_zy, sel_cos_t), zy.get_shape()))
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=updated_logits))
predict_cls = tf.argmax(updated_logits, 1)
accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.cast(predict_cls, tf.int64), tf.cast(labels, tf.int64)), 'float'))
predict_cls_s = tf.argmax(zy, 1)
accuracy_s = tf.reduce_mean(tf.cast(tf.equal(tf.cast(predict_cls_s, tf.int64), tf.cast(labels, tf.int64)), 'float'))
return zy, loss, accuracy, accuracy_s, predict_cls_s