-
Notifications
You must be signed in to change notification settings - Fork 7
/
VGG_Models.py
95 lines (78 loc) · 3.22 KB
/
VGG_Models.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
from torch import nn
cfg = {
'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}
class VGG(nn.Module):
def __init__(self,vgg_name,num_classes=10):
super(VGG, self).__init__()
self.features = self._make_layers(cfg[vgg_name])
self.classifier = nn.Linear(512, num_classes)
def forward(self, x):
out = self.features(x)
out = out.view(out.size(0), -1)
out = self.classifier(out)
return out
def _make_layers(self, cfg):
layers = []
in_channels = 3
for x in cfg:
if x == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
nn.BatchNorm2d(x),
nn.ReLU(inplace=True)]
in_channels = x
layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
return nn.Sequential(*layers)
class VGG_Intermediate_Branches(nn.Module):
def __init__(self,vgg_name,num_classes=10):
super(VGG_Intermediate_Branches, self).__init__()
self.vgg_name= vgg_name
features = self._make_layers(cfg[vgg_name])
if self.vgg_name == "VGG16":
self.features_1 = features[:3]
self.features_2 = features[3:6]
self.features_3 = features[6:10]
self.features_4 = features[10:14]
self.features_5 = features[14:]
else:
self.features_1 = features[:4]
self.features_2 = features[4:7]
self.features_3 = features[7:10]
self.features_4 = features[10:]
self.classifier = nn.Linear(512, num_classes)
def forward(self, x):
intermediate_outputs_list = []
out_1 = self.features_1(x)
intermediate_outputs_list.append(out_1)
out_2 = self.features_2(out_1)
intermediate_outputs_list.append(out_2)
out_3 = self.features_3(out_2)
intermediate_outputs_list.append(out_3)
out_4 = self.features_4(out_3)
intermediate_outputs_list.append(out_4)
if self.vgg_name == "VGG16":
out_5 = self.features_5(out_4)
out = out_5.view(out_5.size(0), -1)
intermediate_outputs_list.append(out_5)
else:
out = out_4.view(out_4.size(0), -1)
out = self.classifier(out)
return out, intermediate_outputs_list
def _make_layers(self, cfg):
layers = []
in_channels = 3
for x in cfg:
if x == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
nn.BatchNorm2d(x),
nn.ReLU(inplace=True)]
in_channels = x
layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
return nn.Sequential(*layers)