VGG有6种子模型,分别是A、A-LRN、B、C、D、E,我们常看到的基本是D、E这两种模型,即VGG16,VGG19
为了方便阅读,并没有加上激活函数层
from torch import nn import torch from torchsummary import summary class VGG16(nn.Module): def __init__(self): super(VGG16, self).__init__() self.sum_Module = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), nn.MaxPool2d(2, 2), nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), nn.MaxPool2d(2, 2), nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1), nn.MaxPool2d(2, 2), nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1), nn.MaxPool2d(2, 2), nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1), nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1), nn.MaxPool2d(2, 2), nn.Flatten(), nn.Linear(7 * 7*512, 4096), #nn.Dropout(0.5), nn.Linear(4096, 4096), #nn.Dropout(0.5), nn.Linear(4096, 1000) ) def forward(self, x): x = self.sum_Module(x), return x if __name__ == '__main__': YOLO = VGG16() device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') inputs = YOLO.to(device) summary(inputs, (3, 224, 224),batch_size=1, device="cuda") # 分别是输入数据的三个维度