VGG model
| 14 | |
| 15 | |
| 16 | class VGG(nn.Module): |
| 17 | ''' VGG model ''' |
| 18 | def __init__(self, features, out_channels): |
| 19 | super(VGG, self).__init__() |
| 20 | self.features = features |
| 21 | self.avgpool = nn.AdaptiveAvgPool2d((1,1)) |
| 22 | self.classifier = nn.Sequential( |
| 23 | nn.Dropout(), |
| 24 | nn.Linear(512, 512), |
| 25 | nn.ReLU(True), |
| 26 | nn.Dropout(), |
| 27 | nn.Linear(512, 512), |
| 28 | nn.ReLU(True), |
| 29 | nn.Linear(512, out_channels), |
| 30 | ) |
| 31 | |
| 32 | ''' Initialize weights ''' |
| 33 | for m in self.modules(): |
| 34 | if isinstance(m, nn.Conv2d): |
| 35 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
| 36 | m.weight.data.normal_(0, math.sqrt(2. / n)) |
| 37 | m.bias.data.zero_() |
| 38 | |
| 39 | |
| 40 | def forward(self, x): |
| 41 | x = self.features(x) |
| 42 | x = self.avgpool(x) |
| 43 | x = torch.flatten(x, 1) |
| 44 | x = self.classifier(x) |
| 45 | return x |
| 46 | |
| 47 | |
| 48 | def make_layers(cfg, in_dims=3, batch_norm=False): |