mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
import mxnet as mx
|
|
from mxnet import autograd
|
|
from mxnet import nd
|
|
|
|
|
|
class MarginLoss(object):
|
|
""" Default is Arcface loss
|
|
"""
|
|
def __init__(self, margins=(1.0, 0.5, 0.0), loss_s=64, embedding_size=512):
|
|
"""
|
|
"""
|
|
# margins
|
|
self.loss_m1 = margins[0]
|
|
self.loss_m2 = margins[1]
|
|
self.loss_m3 = margins[2]
|
|
self.loss_s = loss_s
|
|
self.embedding_size = embedding_size
|
|
|
|
def forward(self, data, weight, mapping_label, depth):
|
|
"""
|
|
"""
|
|
with autograd.record():
|
|
norm_data = nd.L2Normalization(data)
|
|
norm_weight = nd.L2Normalization(weight)
|
|
#
|
|
fc7 = nd.dot(norm_data, norm_weight, transpose_b=True)
|
|
#
|
|
mapping_label_onehot = mx.nd.one_hot(indices=mapping_label,
|
|
depth=depth,
|
|
on_value=1.0,
|
|
off_value=0.0)
|
|
# cosface
|
|
if self.loss_m1 == 1.0 and self.loss_m2 == 0.0:
|
|
_one_hot = mapping_label_onehot * self.loss_m3
|
|
fc7 = fc7 - _one_hot
|
|
else:
|
|
fc7_onehot = fc7 * mapping_label_onehot
|
|
cos_t = fc7_onehot
|
|
t = nd.arccos(cos_t)
|
|
if self.loss_m1 != 1.0:
|
|
t = t * self.loss_m1
|
|
if self.loss_m2 != 0.0:
|
|
t = t + self.loss_m2
|
|
margin_cos = nd.cos(t)
|
|
if self.loss_m3 != 0.0:
|
|
margin_cos = margin_cos - self.loss_m3
|
|
margin_fc7 = margin_cos
|
|
margin_fc7_onehot = margin_fc7 * mapping_label_onehot
|
|
diff = margin_fc7_onehot - fc7_onehot
|
|
fc7 = fc7 + diff
|
|
|
|
fc7 = fc7 * self.loss_s
|
|
return fc7, mapping_label_onehot
|