File size: 2,052 Bytes
998bb30 |
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 |
"""VGG11/13/16/19 in Pytorch."""
import torch.nn as 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.upper()])
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)
def vgg11(num_classes: int) -> VGG:
return VGG('vgg11', num_classes=num_classes)
def vgg13(num_classes: int) -> VGG:
return VGG('vgg13', num_classes=num_classes)
def vgg16(num_classes: int) -> VGG:
return VGG('vgg16', num_classes=num_classes)
def vgg19(num_classes: int) -> VGG:
return VGG('vgg19', num_classes=num_classes)
|