class VGG11(nn.Module):
def __init__(self):
super(VGG11, self).__init__()
self.conv_block1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.conv_block2 = nn.Sequential(
nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.conv_block3 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.conv_block4 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.conv_block5 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.linear1 = nn.Sequential(
nn.Linear(in_features=3136, out_features=512),
nn.ReLU(),
nn.Dropout(p=0.5)
)
self.linear2 = nn.Sequential(
nn.Linear(in_features=512, out_features=512),
nn.ReLU(),
nn.Dropout(p=0.5)
)
self.linear3 = nn.Linear(in_features=512, out_features=120)
def forward(self, x):
x = self.conv_block1(x)
x = self.conv_block2(x)
x = self.conv_block3(x)
x = self.conv_block4(x)
x = self.conv_block5(x)
x = x.view(x.shape[0], -1)
x = self.linear1(x)
x = self.linear2(x)
x = self.linear3(x)
return x
net = VGG11()
summary(net, input_size=(1, 224, 224))
输出:
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 8, 224, 224] 80
ReLU-2 [-1, 8, 224, 224] 0
MaxPool2d-3 [-1, 8, 112, 112] 0
Conv2d-4 [-1, 16, 112, 112] 1,168
ReLU-5 [-1, 16, 112, 112] 0
MaxPool2d-6 [-1, 16, 56, 56] 0
Conv2d-7 [-1, 32, 56, 56] 4,640
ReLU-8 [-1, 32, 56, 56] 0
Conv2d-9 [-1, 32, 56, 56] 9,248
ReLU-10 [-1, 32, 56, 56] 0
MaxPool2d-11 [-1, 32, 28, 28] 0
Conv2d-12 [-1, 64, 28, 28] 18,496
ReLU-13 [-1, 64, 28, 28] 0
Conv2d-14 [-1, 64, 28, 28] 36,928
ReLU-15 [-1, 64, 28, 28] 0
MaxPool2d-16 [-1, 64, 14, 14] 0
Conv2d-17 [-1, 64, 14, 14] 36,928
ReLU-18 [-1, 64, 14, 14] 0
Conv2d-19 [-1, 64, 14, 14] 36,928
ReLU-20 [-1, 64, 14, 14] 0
MaxPool2d-21 [-1, 64, 7, 7] 0
Linear-22 [-1, 512] 1,606,144
ReLU-23 [-1, 512] 0
Dropout-24 [-1, 512] 0
Linear-25 [-1, 512] 262,656
ReLU-26 [-1, 512] 0
Dropout-27 [-1, 512] 0
Linear-28 [-1, 120] 61,560
================================================================
Total params: 2,074,776
Trainable params: 2,074,776
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.19
Forward/backward pass size (MB): 15.65
Params size (MB): 7.91
Estimated Total Size (MB): 23.75