LeNet-5 卷积神经网络模型代码解析

上述代码定义了一个经典的卷积神经网络模型LeNet-5，它包含四个层次的结构：两个卷积层、一个全连接层和一个输出层。

在类的初始化函数中，首先调用了父类的初始化函数，然后定义了四个层次，分别是：

第一层：一个卷积层，包括一个5x5的卷积核，6个输出通道，padding为2，Tanh激活函数和一个2x2的平均池化层。
第二层：一个卷积层，包括一个5x5的卷积核，16个输出通道，padding为0，Tanh激活函数和一个2x2的平均池化层。
第三层：一个卷积层，包括一个5x5的卷积核，120个输出通道。
第四层：一个全连接层，包括一个输入特征数为120，输出特征数为84的线性层，Tanh激活函数，和一个输出特征数为10的线性层。

在模型的前向传播函数中，首先将输入x经过第一层、第二层和第三层的卷积操作，接着将其展平，然后经过第四层的全连接操作，最终输出预测结果。整个模型的输入为一个单通道的28x28的图像，输出为10类的标签。最后将模型移动到指定的设备（GPU或CPU）上。

class LeNet5(nn.Module):
    
    def __init__(self, in_channel, output):
        super(LeNet5, self).__init__()
        self.layer1 = nn.Sequential(nn.Conv2d(in_channels=in_channel, out_channels=6, kernel_size=5, stride=1, padding=2),   # (6, 28, 28)
                                    nn.Tanh(),
                                    nn.AvgPool2d(kernel_size=2, stride=2, padding=0))   # (6, 14, 14))

        self.layer2 = nn.Sequential(nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0),  # (16, 10, 10)
                                    nn.Tanh(),
                                    nn.AvgPool2d(kernel_size=2, stride=2, padding=0))   # (16, 5, 5)

        self.layer3 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5)  # (120, 1, 1)

        self.layer4 = nn.Sequential(nn.Linear(in_features=120, out_features=84),
                                    nn.Tanh(),
                                    nn.Linear(in_features=84, out_features=output))

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = torch.flatten(input=x, start_dim=1)
        x = self.layer4(x)
        return x

model = LeNet5(1, 10)
model.to(device)