Arc-Softmax代码实现(Pytorch)

第一种方法：直接增大θ角，再利用反三角函数

import torch
import torch.nn as nn
import torch.nn.functional as F
class ArcSoftmax(nn.Module):
    def __init__(self,feature_num,cls_num):
         super().__init__()
         self.W=nn.Parameter(torch.randn(feature_num,cls_num))
    def forward(self, feature):
        _W = F.normalize(self.W, dim=0) # 得到W/W模，而norm(W)是的得到W的模
        _X = F.normalize(feature, dim=1)
        cosine = torch.matmul(_X, _W)/10
        s=1
        a = torch.acos(cosine*0.999)
        top = torch.exp(s* torch.cos(a + 1)*10)
        _top = torch.exp(s * cosine*10)
        bottom = torch.sum(torch.exp(s*cosine*10), dim=1,keepdim=True)

        return top / (bottom - _top + top)

第二种方法：利用余弦函数增大θ角

import torch
import torch.nn as nn
import torch.nn.functional as F
class ArcSoftmax(nn.Module):
    def __init__(self,feature_num,cls_num):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(feature_num, cls_num,))
        #  cls_num训练集中总的人脸分类数
        #  feature_num特征向量长度
        self.m = torch.tensor(0.5)
        # 夹角差值 0.5 
        self.cos_m = torch.cos(self.m)
        self.sin_m = torch.sin(self.m)
        # 差值的cos和sin
    def forward(self, feature, label):
        x = F.normalize(feature,dim=1)
        W = F.normalize(self.weight,dim=0).permute(1,0)
        cosine = F.linear(x, W)
        sine = torch.sqrt(1.001-cosine**2)
        phi = cosine * self.cos_m - sine * self.sin_m
        phi_ = torch.where(cosine>0, phi, cosine)
        one_hot = torch.zeros(cosine.size())
        one_hot.scatter_(1, label.view(-1, 1), 1)
        # 将样本的标签映射为one hot形式 例如N个标签，映射为（N，num_classes）
        output = (one_hot * phi_) + ((1.0 - one_hot) * cosine)
        # 对于正确类别 1*cos(θ+β) 即公式中的 cos(θ+β)，对于错误的类别1*cosθ 即公式中的 cosθ
        # 这样对于每一个样本，比如[0,0,0,1,0,0]属于第四类，则最终结果为[cosθ, cosθ, cosθ,cos(θ+β) , cosθ, cosθ]
        # 再乘以半径，经过交叉熵，正好是ArcFace的公式

        return output*64,cosine