mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
67 lines
1.6 KiB
Python
67 lines
1.6 KiB
Python
import oneflow as flow
|
|
from oneflow import nn
|
|
|
|
|
|
def get_loss(name):
|
|
if name == "cosface":
|
|
return CosFace()
|
|
elif name == "arcface":
|
|
return ArcFace()
|
|
else:
|
|
raise ValueError()
|
|
|
|
|
|
class CrossEntropyLoss_sbp(nn.Module):
|
|
def __init__(self):
|
|
super(CrossEntropyLoss_sbp, self).__init__()
|
|
|
|
def forward(self, logits, label):
|
|
loss = flow._C.sparse_softmax_cross_entropy(
|
|
logits, label)
|
|
loss = flow.mean(loss)
|
|
return loss
|
|
|
|
|
|
def get_loss(name):
|
|
if name == "cosface":
|
|
return CosFace()
|
|
elif name == "arcface":
|
|
return ArcFace()
|
|
else:
|
|
raise ValueError()
|
|
|
|
|
|
class CosFace(nn.Module):
|
|
def __init__(self, s=64.0, m=0.40):
|
|
super(CosFace, self).__init__()
|
|
self.s = s
|
|
self.m = m
|
|
|
|
def forward(self, cosine, label):
|
|
index = flow.where(label != -1)[0]
|
|
m_hot = flow.zeros(index.size()[0], cosine.size()[
|
|
1], device=cosine.device)
|
|
|
|
m_hot = flow.scatter(m_hot, 1, label[index, None], self.m)
|
|
cosine = cosine[index] - m_hot
|
|
|
|
ret = cosine * self.s
|
|
return ret
|
|
|
|
|
|
class ArcFace(nn.Module):
|
|
def __init__(self, s=64.0, m=0.5):
|
|
super(ArcFace, self).__init__()
|
|
self.s = s
|
|
self.m = m
|
|
|
|
def forward(self, cosine: flow.Tensor, label):
|
|
index = flow.where(label != -1)[0]
|
|
m_hot = flow.zeros(index.size()[0], cosine.size()[
|
|
1], device=cosine.device)
|
|
m_hot.scatter_(1, label[index, None], self.m)
|
|
cosine.acos_()
|
|
cosine[index] += m_hot
|
|
cosine.cos_().mul_(self.s)
|
|
return cosine
|