本文最后更新于:14 天前
示例代码:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2018/9/14 13:33
# @Author : Seven
# @Site :
# @File : VGG.py
# @Software: PyCharm
import torch.nn as nn
model = {
'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):
super(VGG, self).__init__()
self.features = self._make_layers(model[vgg_name])
self.classifier = nn.Linear(512, 10)
def forward(self, x):
network = self.features(x)
network = network.view(network.size(0), -1)
out = self.classifier(network)
return out, network
@staticmethod
def _make_layers(models):
layers = []
in_channels = 3
for layer in models:
if layer == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
layers += [nn.Conv2d(in_channels, layer, kernel_size=3, padding=1),
nn.BatchNorm2d(layer),
nn.ReLU(inplace=True)]
in_channels = layer
layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
return nn.Sequential(*layers)
本博客所有文章除特别声明外,均采用 CC BY-SA 3.0协议 。转载请注明出处!