Open tfygg opened 3 years ago
`class ArcMarginProduct(nn.Module): def init(self, in_feature=128, out_feature=10575, s=32.0, m=0.50, easy_margin=False): super(ArcMarginProduct, self).init() self.in_feature = in_feature self.out_feature = out_feature self.s = s self.m = m self.weight = Parameter(torch.Tensor(out_feature, in_feature)) nn.init.xavieruniform(self.weight)
self.easy_margin = easy_margin self.cos_m = math.cos(m) self.sin_m = math.sin(m) # make the function cos(theta+m) monotonic decreasing while theta in [0°,180°] self.th = math.cos(math.pi - m) self.mm = math.sin(math.pi - m) * m def forward(self, x, label): # cos(theta) cosine = F.linear(F.normalize(x), F.normalize(self.weight)) # cos(theta + m) sine = torch.sqrt(1.0 - torch.pow(cosine, 2)) phi = cosine * self.cos_m - sine * self.sin_m if self.easy_margin: phi = torch.where(cosine > 0, phi, cosine) else: phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm) #one_hot = torch.zeros(cosine.size(), device='cuda' if torch.cuda.is_available() else 'cpu') 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 return output`
This is the implementation of amsoftmax, but in your code, cos_theta is equal to output without normalization.
This function takes cosine from the last layer of the model. See class AngleSimpleLinear.
`class ArcMarginProduct(nn.Module): def init(self, in_feature=128, out_feature=10575, s=32.0, m=0.50, easy_margin=False): super(ArcMarginProduct, self).init() self.in_feature = in_feature self.out_feature = out_feature self.s = s self.m = m self.weight = Parameter(torch.Tensor(out_feature, in_feature)) nn.init.xavieruniform(self.weight)
This is the implementation of amsoftmax, but in your code, cos_theta is equal to output without normalization.