关于pytorch官网教程中的What is torch.nn really?（三）

本文介绍如何使用PyTorch搭建并训练卷积神经网络(CNN)解决MNIST手写数字识别问题，包括模型定义、数据预处理及GPU加速等关键步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

文章目录

因为MNIST是个分类问题，所以使用逻辑回归是合适的，同时又是个图像问题，所以使用CNN也是合理的。
接下来我们转向CNN，这里展示了两种方式，因为之前的代码已经把model封装起来了，而其余部分并没有涉及到模型本身，所以基本上只要重新定义一个模型，其余部分都不需要改动。

Switch to CNN

class Mnist_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1)

    def forward(self, xb):
        xb = xb.view(-1, 1, 28, 28)
        xb = F.relu(self.conv1(xb))
        xb = F.relu(self.conv2(xb))
        xb = F.relu(self.conv3(xb))
        xb = F.avg_pool2d(xb, 4)
        return xb.view(-1, xb.size(1))

首先，当然，模型类的名字从Mnist_Logistic改成了Mnist_CNN，名字的变化已经意味着这里从逻辑回归转向了CNN，或者说卷积神经网络。
其次，还是两个函数，一个__init__，一个forward，当然这两个函数定义与逻辑回归的模型肯定是不一样的。大体上，我们可以把__init__理解为定义了一个model的框架，而forward则是在这个框架上处理数据的过程。
这里用Conv2d定义了一个3层的卷积网络，注意__init__定义中并没有对输入数据做任何假定，只是定义了卷积网络本身的层次。
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)表示这一层的输入channel(in_channels)是1，因为是第一层，输入的是黑白图片，所以in_channels是1，输出channel(out_channels)是16，每次计算 $kernel\_size\times kernel\_size$ 也就是 $3\times3$ 的格子，步长为2，以及需要填充边缘，因为padding_mode为缺省值'zeros'，所以是用0来填充。其余两层的定义是类似的，就不废话了。
至于xb.view(-1, 1, 28, 28)，Tensor.view的文档在这里，就这里而言，是把一个torch.Size([64, 784])的张量变形为一个torch.Size([64, 1, 28, 28])的张量。参数列表中的-1表示这个值由其他参数推断，因为变形前后的数据的数量必须保持一致，所以这里就推算出来是64： $64\times784=64\times 1\times 28\times 28$ ，这个64是每个批次的大小。很容易能够看出来，这里是相当于把本来拉成了一个784维的行向量的 $28 \times 28$ 的图片恢复成了 $28 \times 28$