AlexNet代码解读

AlexNet代码解读

概述

AlexNet的网络结构很简单，是最初级版本的CNN，没有使用什么技巧。
网络分成两个部分，分别是卷积、激活、池化构成的特征提取器，以及前向神经网络的分类器。

网络结构图

AlexNet代码细节分析

import numpy as np
import torch
import torch.nn as nn
from typing import Any
from torchsummary import summary
__all__ = ['AlexNet','alexnet']

mpdel_urls =  {
    'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',
}
class AlexNet(nn.Module):
class AlexNet(nn.Module):
    def __init__(self, num_classes: int = 1000) -> None:
        super(AlexNet, self).__init__()
         # conv、relu、maxpool串联结构来提取特征
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        # 分类器（前面的卷积层已经全部写好，提取出特征了）
        # AlexNet的卷积层比较简单，层数不深，就直接写在features函数里面了
        # 特征层操作：（卷积、激活、池化）*2、（卷积、激活）*2、卷积、激活、池化
        self.avgpool = nn.AdaptiveAvgPool2d(6)
        # 分类器操作：(dropout、全连接、激活)*2、全连接
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x
# 构建alexnet
def alexnet(pretrained:bool = False, progress:bool = True, **kwargs:Any)->AlexNet:
    model = AlexNet(**kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls['alexnet'],progress = progress)
        model.load_state_dict(state_dict)
    return model

现在输入一个3x224x224的tensor，看它经过alexnet每一层之后会变成怎样大小的tensor。

from torchsummary import summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
t = AlexNet().to(device)
summary(t,(3,224,224))

好了，输出结果如下所示。

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 64, 55, 55]          23,296
              ReLU-2           [-1, 64, 55, 55]               0
         MaxPool2d-3           [-1, 64, 27, 27]               0
            Conv2d-4          [-1, 192, 27, 27]         307,392
              ReLU-5          [-1, 192, 27, 27]               0
         MaxPool2d-6          [-1, 192, 13, 13]               0
            Conv2d-7          [-1, 384, 13, 13]         663,936
              ReLU-8          [-1, 384, 13, 13]               0
            Conv2d-9          [-1, 256, 13, 13]         884,992
             ReLU-10          [-1, 256, 13, 13]               0
           Conv2d-11          [-1, 256, 13, 13]         590,080
             ReLU-12          [-1, 256, 13, 13]               0
        MaxPool2d-13            [-1, 256, 6, 6]               0
AdaptiveAvgPool2d-14            [-1, 256, 6, 6]               0
          Dropout-15                 [-1, 9216]               0
           Linear-16                 [-1, 4096]      37,752,832752,832                                                       0
             ReLU-17                 [-1, 4096]               0
      0                                                 781,312
          Dropout-18                 [-1, 4096]               0
      0                                                 097,000
           Linear-19                 [-1, 4096]      16,========781,312
             ReLU-20                 [-1, 4096]
      0
           Linear-21                 [-1, 1000]       4,097,000
================================================================
Total params: 61,100,840
Trainable params: 61,100,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 8.38
Params size (MB): 233.08
Estimated Total Size (MB): 242.03
----------------------------------------------------------------