mccreezhao / fr-loss-on-mnist Goto Github PK

def forward(self, x, label):
    cosine = F.normalize(x).mm(F.normalize(self.weight, dim=0))
    sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
    phi = cosine * self.cos_m - sine * self.sin_m  # 两角和公式
    # 为了保证cos(theta+m)在0-pi单调递减：
    phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm)
    one_hot = torch.zeros_like(cosine)
    one_hot.scatter_(1, label.view(-1, 1), 1)
    output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
    output = output * self.s
    #这里这个s比较关键，如果设置的不够大会导致收敛困难，
    # 比较容易出现两个类别的Feature重叠的现象，
    # 这也是加性Margin相对于乘性Margin的一个比较弱势的地方，
    # 对于特征向量的收缩要求足够，但是对两类特征向量之间的距离约束不够
    loss = F.cross_entropy(cosine, label)

    return loss, output