mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: liaoxingyu
|
|
@contact: sherlockliao01@gmail.com
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.nn import Parameter
|
|
|
|
|
|
class AM_softmax(nn.Module):
|
|
r"""Implement of large margin cosine 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, out_features, s=30.0, m=0.40):
|
|
super().__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.s = s
|
|
self.m = m
|
|
self.weight = Parameter(torch.FloatTensor(out_features, in_features))
|
|
# nn.init.normal_(self.weight, std=0.001)
|
|
nn.init.xavier_uniform_(self.weight)
|
|
|
|
def forward(self, input, label):
|
|
# --------------------------- cos(theta) & phi(theta) ---------------------------
|
|
cosine = F.linear(F.normalize(input), F.normalize(self.weight)) # (bs, num_classes)
|
|
phi = cosine - self.m
|
|
# phi = cosine
|
|
# --------------------------- convert label to one-hot ---------------------------
|
|
one_hot = torch.zeros(cosine.size()).to(label.device)
|
|
# one_hot = one_hot.cuda() if cosine.is_cuda else one_hot
|
|
one_hot.scatter_(1, label.view(-1, 1).long(), 1)
|
|
# -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
|
|
# you can use torch.where if your torch.__version__ is 0.4
|
|
output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
|
|
# output *= torch.norm(input, p=2, dim=1, keepdim=True)
|
|
output *= self.s
|
|
|
|
return output
|
|
|