ArcFace与 AM-Softmax同样也是加性的 margin，差别只是 ArcFace的 margin 加在 Cos 算子的里面，而 AM-Softmax的 margin 在加性算子的外面。 ArcFace的loss伪代码如下：

    • 对x进行归一化
    • 对W进行归一化
    • 计算Wx得到预测向量y
    • 从y中挑出与ground truth对应的值
    • 计算其反余弦得到角度
    • 角度加上m
    • 从y 中挑出与ground truth对应的值所在位置的独热码
    • 将 cos(theta+m)通过独热码放回原来的位置
    • 对所有值乘上尺度因子S
    1.  ArcFace
    2. class ArcMarginProduct(nn.Module):
    3.     r"""Implement of large margin arc distance: :
    4.         Args:
    5.             in_features: size of each input sample
    6.             out_features: size of each output sample
    7.             s: norm of input feature
    8.             m: margin
    10.             cos(theta + m)
    11.         """
    13.     def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False):
    14.         super(ArcMarginProduct, self).__init__()
    15.         self.in_features = in_features
    16.         self.out_features = out_features
    17.         self.s = s
    18.         self.m = m
    19.         # Parameter 的用途：
    20.         # 将一个不可训练的类型Tensor转换成可以训练的类型parameter
    21.         # 并将这个parameter绑定到这个module里面
    22.         # net.parameter()中就有这个绑定的parameter，所以在参数优化的时候可以进行优化的
    23.         # https://www.jianshu.com/p/d8b77cc02410
    24.         # 初始化权重
    25.         self.weight = Parameter(torch.FloatTensor(out_features, in_features))
    26.         nn.init.xavier_uniform_(self.weight)
    28.         self.easy_margin = easy_margin
    29.         self.cos_m = math.cos(m)
    30.         self.sin_m = math.sin(m)
    31.         self.th = math.cos(math.pi - m)
    32.         self.mm = math.sin(math.pi - m) * m
    34.     def forward(self, input, label):
    35.         # --------------------------- cos(theta) & phi(theta) ---------------------------
    36.         # torch.nn.functional.linear(input, weight, bias=None)
    37.         # y=x*W^T+b
    38.         cosine = F.linear(F.normalize(input), F.normalize(self.weight))
    39.         sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
    40.         # cos(a+b)=cos(a)*cos(b)-size(a)*sin(b)
    41.         phi = cosine * self.cos_m - sine * self.sin_m
    42.         if self.easy_margin:
    43.             # torch.where(condition, x, y) → Tensor
    44.             # condition (ByteTensor) – When True (nonzero), yield x, otherwise yield y
    45.             # x (Tensor) – values selected at indices where condition is True
    46.             # y (Tensor) – values selected at indices where condition is False
    47.             # return:
    48.             # A tensor of shape equal to the broadcasted shape of condition, x, y
    49.             # cosine>0 means two class is similar, thus use the phi which make it
    50.             phi = torch.where(cosine > 0, phi, cosine)
    51.         else:
    52.             phi = torch.where(cosine > self.th, phi, cosine - self.mm)
    53.         # --------------------------- convert label to one-hot ---------------------------
    54.         # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')
    55.         # 将cos(\theta + m)更新到tensor相应的位置中
    56.         one_hot = torch.zeros(cosine.size(), device='cuda')
    57.         # scatter_(dim, index, src)
    58.         one_hot.scatter_(1, label.view(-1, 1).long(), 1)
    59.         # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
    60.         output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
    61.         # you can use torch.where if your torch.__version__ is 0.4
    62.         output *= self.s
    63.         # print(output)
    65.         return output