-
Notifications
You must be signed in to change notification settings - Fork 20
/
Copy pathmodel.py
123 lines (106 loc) · 4.33 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# part of this code borrows from https://github.com/layumi/Person_reID_baseline_pytorch && https://github.com/ZhaoJ9014/face.evoLVe.PyTorch/blob/master/head/metrics.py
import torch
import torch.nn as nn
from torchvision import models
from torch.autograd import Variable
class ft_net(nn.Module):
def __init__(self, feature_dim, num_classes, num_gpus=1, am=False, model_parallel=False, class_split=None):
super(ft_net, self).__init__()
model_ft = models.resnet50(pretrained=True)
model_ft.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.model = model_ft.cuda()
self.backbone = nn.Sequential(
model_ft.conv1,
model_ft.bn1,
model_ft.relu,
model_ft.maxpool,
model_ft.layer1,
model_ft.layer2,
model_ft.layer3,
model_ft.layer4,
model_ft.avgpool
).cuda()
self.features = nn.Linear(2048, feature_dim).cuda()
if am:
self.classifier = FullyConnected_AM(feature_dim, num_classes, num_gpus, model_parallel, class_split)
else:
self.classifier = FullyConnected(feature_dim, num_classes, num_gpus, model_parallel, class_split)
def forward(self, x, labels=None):
x = self.backbone(x)
x = x.view(x.size(0), x.size(1))
x = self.features(x)
x = self.classifier(x, labels)
return x
class FullyConnected(nn.Module):
def __init__(self, in_dim, out_dim, num_gpus=1, model_parallel=False, class_split=None):
super(FullyConnected, self).__init__()
self.num_gpus = num_gpus
self.model_parallel = model_parallel
if model_parallel:
self.fc_chunks = nn.ModuleList()
for i in range(num_gpus):
self.fc_chunks.append(
nn.Linear(in_dim, class_split[i], bias=False).cuda(i)
)
else:
self.fc = nn.Linear(in_dim, out_dim, bias=False)
def forward(self, x, labels=None):
if self.model_parallel:
x_list = []
for i in range(self.num_gpus):
_x = self.fc_chunks[i](x.cuda(i))
x_list.append(_x)
return tuple(x_list)
else:
return self.fc(x)
class FullyConnected_AM(nn.Module):
def __init__(self, in_dim, out_dim, num_gpus=1, model_parallel=False, class_split=None, margin=0.35, scale=30):
super(FullyConnected_AM, self).__init__()
self.num_gpus = num_gpus
self.model_parallel = model_parallel
if self.model_parallel:
self.am_branches = nn.ModuleList()
for i in range(num_gpus):
self.am_branches.append(AM_Branch(in_dim, class_split[i], margin, scale).cuda(i))
else:
self.am = AM_Branch(in_dim, out_dim, margin, scale)
def forward(self, x, labels=None):
if self.model_parallel:
output_list = []
for i in range(self.num_gpus):
output = self.am_branches[i](x.cuda(i), labels[i])
output_list.append(output)
return tuple(output_list)
else:
return self.am(x, labels)
class AM_Branch(nn.Module):
def __init__(self, in_dim, out_dim, margin=0.35, scale=30):
super(AM_Branch, self).__init__()
self.m = margin
self.s = scale
# training parameter
self.weight = nn.Parameter(torch.Tensor(in_dim, out_dim), requires_grad=True)
self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
def forward(self, x, label):
x_norm = x.pow(2).sum(1).pow(0.5)
w_norm = self.weight.pow(2).sum(0).pow(0.5)
cos_theta = torch.mm(x, self.weight) / x_norm.view(-1, 1) / w_norm.view(1, -1)
cos_theta = cos_theta.clamp(-1, 1)
phi = cos_theta - self.m
index = label.data
index = index.byte()
output = cos_theta * 1.0
output[index] = phi[index]
output *= self.s
return output
if __name__ == '__main__':
net = ft_net(256, 65536, 4, False)
print(net)
input = Variable(torch.FloatTensor(8, 3, 256, 128))
output = net(input)
print('net output size:')
if isinstance(output, tuple):
for o in output:
print(o.shape)
else:
print(output.shape)