cosine face 的pytorch实现

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
import math

class ArcMargin(nn.Module):
    r"""Implement of large margin arc distance: :
        Args:
          in_features: size of each input sample
          out_features: size of each output sample
          s: norm of input feature
          m: margin
          cos(theta + m)
      """
    def __init__(self, in_features=2048, out_features=2000, s=64, m=0.45):
        super(ArcMargin, self).__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.weight = Parameter(torch.cuda.FloatTensor(out_features, in_features))
        nn.init.normal_(self.weight, std=0.001)
        #nn.init.xavier_uniform_(self.weight)
        self.classifier = nn.Linear(in_features, out_features, bias=False)  
        self.classifier.apply(weights_init_classifier)

    def forward(self, inputs, targets):
        # --------------------------- cos(theta) & phi(theta) ---------------------------
        cosine = F.linear(F.normalize(inputs, p=2), F.normalize(self.weight, p=2))
        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')
        one_hot = torch.zeros(cosine.size(), device='cuda').scatter_(1, targets.view(targets.size(0), 1).long(), s_m)
        cos_feat = self.s * cosine - one_hot
        loss = torch.nn.functional.cross_entropy(cos_feat, targets) 

        # print(loss)
        return loss