CNN 卷积神经网络 （MNIST 手写数字数据集的分类）

一、MNIST 手写数字数据集的分类

1.数据集加载与可视化

# 加载MNIST训练集和测试集
train_data=datasets.MNIST(
    root="data",        # 数据保存路径
    train=True,         # 加载训练集
    download=True,      # 如果本地没有则自动下载
    transform=ToTensor(),# 将图片转为PyTorch的Tensor格式，并归一化到[0,1]
)

test_data=datasets.MNIST(
    root="data",
    train=False,        # 加载测试集
    download=True,
    transform=ToTensor(),
)

# 可视化9张样本图片（从第59000个样本开始）
from matplotlib import pyplot as plt
figure=plt.figure()
for i in range(9):
    img,label=train_data[i+59000]  # 获取第59000+i个样本的图片和标签

    figure.add_subplot(3,3,i+1)    # 创建3x3的子图布局
    plt.title(label)               # 显示图片对应的数字标签
    plt.axis("off")                # 关闭坐标轴
    plt.imshow(img.squeeze(),cmap="gray")  # 显示图片（squeeze()去掉维度为1的通道维度）
    a=img.squeeze()  # 只是临时保存处理后的图片，无实际作用

plt.show()  # 展示图片

ToTensor()：将 PIL 图片转为形状为[C, H, W]（通道、高度、宽度）的 Tensor，且像素值从[0,255]归一化到[0,1]

img.squeeze()：MNIST 图片是单通道（维度为 1），squeeze()会去掉这个维度，从[1,28,28]变为[28,28]，才能被 matplotlib 正确显示

2. 数据加载器（DataLoader）创建

# 创建数据加载器，按批次加载数据
train_dataloader=DataLoader(train_data,batch_size=64)  # 训练集批次大小64
test_dataloader=DataLoader(test_data,batch_size=64)    # 测试集批次大小64

DataLoader：PyTorch 的核心数据加载工具，将数据集按batch_size分成多个批次，支持多线程加载、打乱数据等

3. 设备配置

device ="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

1.这是 PyTorch 跨平台的设备配置方式，确保代码能在不同硬件上运行

2.后续模型和数据都会被移到这个设备上，利用 GPU 加速计算

4. CNN 模型定义

class CNN(nn.Module):#通过调用类的形式来使用神经网络，神经网络的模型，nn.module1个用法
    def __init__(self):#python基础关于类，self类自己本身
        super(CNN,self).__init__()#继承的父类初始化
        self.conv1=nn.Sequential(
            nn.Conv2d(1,16,5,1,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.MaxPool2d(kernel_size=2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.ReLU(),
        )
        # self.conv3=nn.Sequential(
        #     nn.Conv2d(32, 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)#把刚刚创建的模型传入到Gpu
print(model)

1）.nn.Sequential：按顺序堆叠网络层，简化代码

2）.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)：卷积层参数说明：

.padding=2：在图片边缘填充 2 个像素，保证卷积后尺寸不变（计算公式：H_out = (H_in + 2*padding - kernel_size)/stride + 1）

4）.x.view(x.size(0),-1)：view是重塑张量形状，-1表示自动计算该维度大小，目的是将二维特征图转为一维向量，输入全连接层

5）.model.to(device)：将模型的所有参数移到指定设备

5.训练函数

def train(dataloader,model,loss_fn,optimizer):
    model.train()  # 切换到训练模式（启用Dropout、BatchNorm等训练特有的层）
    batch_size_num=1  # 统计批次数量
    
    for x,y in dataloader:  # 遍历每个批次
        x,y = x.to(device),y.to(device)  # 数据移到指定设备
        pred= model.forward(x)           # 前向传播得到预测值（可简写为model(x)）
        loss = loss_fn(pred,y)           # 计算预测值和真实标签的损失
        
        # 反向传播三部曲
        optimizer.zero_grad()  # 清空上一轮的梯度（否则会累加）
        loss.backward()        # 反向传播计算梯度
        optimizer.step()       # 根据梯度更新模型参数
        
        loss_value=loss.item() # 提取损失值（从Tensor转为普通数值）
        if batch_size_num %100 ==0:  # 每100个批次打印一次损失
            print(f"loss:{loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num +=1

• model.train()：启用训练模式，比如 BatchNorm 会更新均值 / 方差，Dropout 会随机丢弃神经元

• optimizer.zero_grad()：PyTorch 的梯度会累加，所以每次迭代前必须清空

• loss.backward()：计算所有可训练参数的梯度

• optimizer.step()：用优化器更新参数（Adam 算法）

6.测试函数

def test(dataloader,model,loss_fn):
    size =len(dataloader.dataset)  # 测试集总样本数（10000）
    num_batches=len(dataloader)    # 测试集总批次
    model.eval()                   # 切换到评估模式（关闭Dropout、固定BatchNorm）
    test_loss,correct =0,0         # 初始化总损失和正确数
    
    with torch.no_grad():  # 禁用梯度计算（测试时不需要，节省内存和计算）
        for x,y in dataloader:
            x,y = x.to(device),y.to(device)
            pred=model.forward(x)
            test_loss+=loss_fn(pred,y).item()  # 累加批次损失
            # 计算正确预测数：pred.argmax(1)取每行最大值索引（预测类别），和真实标签比较
            correct +=(pred.argmax(1)== y).type(torch.float).sum().item()
    
    test_loss/=num_batches  # 计算平均损失
    correct/=size           # 计算准确率
    print(f"Test result:\n Accuracy:{(100*correct)}%，Avg loss: {test_loss}")

• model.eval()：切换到评估模式，保证 Dropout/BatchNorm 的行为和训练一致

• torch.no_grad()：上下文管理器，禁用梯度计算，大幅提升测试速度

• pred.argmax(1)：在维度 1 上取最大值索引（每个样本的预测类别）

7.测试结果

# 定义损失函数（交叉熵损失，适合分类任务）
loss_fn=nn.CrossEntropyLoss()

# 定义优化器（Adam，学习率0.01）
optimizer=torch.optim.Adam(model.parameters(),lr=0.01)

# 训练10个epoch（完整遍历数据集10次）
for t in range(10):
    print(f"Epoch {t+1}\n----------")
    train(train_dataloader,model,loss_fn,optimizer)
print("Done!")

# 训练完成后测试模型
test(test_dataloader,model,loss_fn)

nn.CrossEntropyLoss()：结合了log_softmax和nll_loss，是分类任务的标准损失函数 torch.optim.Adam：Adam 优化器是目前最常用的优化器之一，自适应学习率