LeNet-5的pytorch实现

from torch import nn
import torch
from torch.nn import functional as F

class LeNet(nn.Module):
    def __init__(self, num_class = 10):
        #num_class为需要分到的类别数
        super().__init__()
        #输入像素大小为1*28*28
        self.features = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size = 5, padding = 2),#输出为6*28*28
            nn.AvgPool2d(kernel_size= 2, stride= 2),#输出为6*14*14，此处也可用MaxPool2d
            nn.Conv2d(6, 16, kernel_size = 5),#输出为16*10*10
            nn.ReLU(),#论文中为sigmoid，但极易出现梯度消失
            nn.AvgPool2d(kernel_size= 2, stride= 2),#输出为16*5*5
            nn.Flatten()#将通道及像素进行合并，方便进一步使用全连接层
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(), #论文中同样为sigmoid
            nn.Linear(120, 84),
            nn.Linear(84, 10))
    def forward(self, x):
            x = self.features(x)
            x = self.classfier(x)

网络结构：

LeNet(
  (features): Sequential(
    (0): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): AvgPool2d(kernel_size=2, stride=2, padding=0)
    (2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
    (3): ReLU()
    (4): AvgPool2d(kernel_size=2, stride=2, padding=0)
    (5): Flatten(start_dim=1, end_dim=-1)
  )
  (classifier): Sequential(
    (0): Linear(in_features=400, out_features=120, bias=True)
    (1): ReLU()
    (2): Linear(in_features=120, out_features=84, bias=True)
    (3): Linear(in_features=84, out_features=10, bias=True)
  )
)