一、准备工作

导入库，导入数据集，划分训练批次数量，规定训练硬件（这部分

import torch
from torch import nn  # 导入神经网络模块
from torch.utils.data import DataLoader  # 数据包管理工具，打包数据
from torchvision import datasets  # 封装了很多与图像相关的模型，和数据集
from torchvision.transforms import ToTensor  # 将其他数据类型转化为张量train_data = datasets.MNIST(root='data',train=True,  # 是否读取下载后数据中的训练集download=True,  # 如果之前下载过则不用下载transform=ToTensor()
)
test_data = datasets.MNIST(root='data',train=False,download=True,transform=ToTensor()
)train_dataloader = DataLoader(train_data,batch_size=256)#是一个类，现在初始化了，但没开始打包，训练开始才打包
test_dataloader = DataLoader(test_data,batch_size=256)device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
print(f'Using {device} device')

二、定义神经网络（重点）

这部分比较重要，我分开讲

1、类定义与继承

class CNN(nn.Module):

这里定义了一个名为CNN的类，它继承自 PyTorch 的nn.Module类。nn.Module是 PyTorch 中所有神经网络模块的基类，通过继承它，我们可以利用 PyTorch 提供的各种功能，如参数管理、设备迁移等。

2、初始化方法

def __init__(self):super().__init__()

这是类的构造函数，super().__init__()调用了父类nn.Module的构造函数，确保父类得到正确初始化。

3、网络层定义

• 第一个卷积块（conv1）

self.conv1 = nn.Sequential(nn.Conv2d(in_channels=1,      # 输入通道数，1表示灰度图像out_channels=16,    # 输出通道数/卷积核数量kernel_size=5,      # 卷积核大小5×5stride=1,           # 步长为1padding=2,          # 填充为2，保持特征图大小不变),nn.ReLU(),               # ReLU激活函数nn.MaxPool2d(kernel_size=2),  # 2×2最大池化
)

这个卷积块接收 1 通道的输入，通过 16 个 5×5 的卷积核进行卷积操作，然后经过 ReLU 激活和 2×2 的最大池化。

• 第二个卷积块（conv2）

self.conv2 = nn.Sequential(nn.Conv2d(16,32,5,1,2),  # 16→32通道，5×5卷积核nn.ReLU(),nn.Conv2d(32,32,5,1,2),  # 32→32通道，5×5卷积核nn.ReLU(),nn.Conv2d(32,32,5,1,2),  # 32→32通道，5×5卷积核nn.ReLU(),nn.Conv2d(32,64,5,1,2),  # 32→64通道，5×5卷积核nn.ReLU(),nn.Conv2d(64,64,5,1,2),  # 64→64通道，5×5卷积核nn.ReLU(),nn.MaxPool2d(kernel_size=2),  # 2×2最大池化
)

这个卷积块包含多个卷积层，逐步增加通道数，并在最后进行一次最大池化。

• 第三个卷积块（conv3）

self.conv3 = nn.Sequential(nn.Conv2d(64,64,5,1,2),  # 64→64通道，5×5卷积核nn.ReLU(),
)

这是一个简单的卷积块，保持通道数不变。

• 全连接层（out）

self.out = nn.Linear(64*7*7,10)

这是网络的输出层，将卷积得到的特征图展平后映射到 10 个输出（可能对应 10 类分类问题）。

4、前向传播方法

def forward(self, x):x = self.conv1(x)    # 通过第一个卷积块x = self.conv2(x)    # 通过第二个卷积块x = self.conv3(x)    # 通过第三个卷积块x = x.view(x.size(0),-1)  # 展平特征图，保留批次维度output = self.out(x)  # 通过全连接层得到输出return output

forward方法定义了数据在网络中的流动路径，即前向传播过程。x.view(x.size(0),-1)将卷积操作得到的多维特征图展平成一维向量，以便输入到全连接层。

5、模型实例化

model = CNN().to(device)

创建 CNN 类的实例，并将模型迁移到指定的设备（CPU 或 GPU）上。

完整代码：

 定义神经网络，通过类的继承
class CNN(nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Sequential(#容器，添加网络层nn.Conv2d(in_channels=1,out_channels = 16,kernel_size = 5,stride = 1,padding = 2,),nn.ReLU(),nn.MaxPool2d(kernel_size = 2),)self.conv2 = nn.Sequential(  # 容器，添加网络层nn.Conv2d(16,32,5,1,2),nn.ReLU(),nn.Conv2d(32, 32, 5, 1, 2),nn.ReLU(),nn.Conv2d(32, 32, 5, 1, 2),nn.ReLU(),nn.Conv2d(32,64,5,1,2),nn.ReLU(),nn.Conv2d(64, 64, 5, 1, 2),nn.ReLU(),nn.MaxPool2d(kernel_size=2),)self.conv3 = nn.Sequential(nn.Conv2d(64,64,5,1,2),nn.ReLU(),)self.out = nn.Linear(64*7*7,10)def forward(self, x):  # 前向传播x = self.conv1(x)x = self.conv2(x)x = self.conv3(x)x = x.view(x.size(0),-1)output = self.out(x)return output
model = CNN().to(device)

三、模型的训练

这一段和上一个博客的步骤一样这里就不多做讲解了

不了解的直接看

深度学习----由手写数字识别案例来认识PyTorch框架-CSDN博客

def train(dataloader, model, loss_fn, optimizer):model.train()  # 开启模型训练模式，像 Dropout、BatchNorm 层会在训练/测试时表现不同，需此设置batch_size_num = 1  # 用于计数当前处理到第几个 batchfor X, y in dataloader:  # 从数据加载器中逐个取出 batch 的数据（特征 X、标签 y ）X, y = X.to(device), y.to(device)  # 把数据和标签放到指定计算设备（CPU/GPU）pred = model(X)  # 将数据输入模型，得到预测结果（模型自动做前向传播计算 ）loss = loss_fn(pred, y)  # 用损失函数计算预测结果和真实标签的损失# 以下是反向传播更新参数的标准流程optimizer.zero_grad()  # 清空优化器里参数的梯度，避免梯度累加影响计算loss.backward()  # 反向传播，计算参数的梯度optimizer.step()  # 根据梯度，更新模型参数loss = loss.item()  # 取出损失张量的数值（脱离计算图 ）# 打印当前 batch 的损失和 batch 编号if batch_size_num % 100 == 0:print(f"loss:{loss:>7f}  [number:{batch_size_num}")batch_size_num += 1  # batch 计数加一def test(dataloader, model, loss_fn):size = len(dataloader.dataset)num_batches = len(dataloader)model.eval()test_loss,correct = 0,0with torch.no_grad():for X , y in dataloader:X ,y = X.to(device),y.to(device)pred = model.forward(X)#.forward可以被省略test_loss += loss_fn(pred,y).item()correct += (pred.argmax(1) == y).type(torch.float).sum().item()a = (pred.argmax(1) == y)b = (pred.argmax(1) == y).type(torch.float)test_loss /= num_batchescorrect /= sizeprint(f"test result: \n Accuracy: {(100*correct)}%,Avg loss:{test_loss}")# print(list(model.parameters()))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.001)#换优化器可以提高准确率，Adam，SGD等# train(train_dataloader,model,loss_fn,optimizer)
# test(test_dataloader,model,loss_fn)
#epochs = 10
for t in range(epochs):print(f"轮次:{t+1}\n----------------------------")train(train_dataloader,model,loss_fn,optimizer)
print("Done")
test(test_dataloader,model,loss_fn)