PyTorch是一个开源的深度学习框架,由 Facebook 的人工智能研究团队开发,广泛应用于计算机视觉(CV)、自然语言处理(NLP)等领域。它以灵活性和易用性著称,特别适合研究人员和开发者进行快速原型设计和实验。

💡PyTorch 的模型训练通常包括以下步骤:

  1. 数据加载:使用 Dataset 和 DataLoader 加载和预处理数据。

  2. 模型定义:通过继承 torch.nn.Module 构建神经网络。

  3. 损失函数与优化器:选择合适的损失函数(如交叉熵)和优化器(如 SGD 或 Adam)。

  4. 训练与验证:通过循环迭代训练模型,并使用验证集评估性能。

  5. 保存与加载模型:使用 torch.save 和 torch.load 保存和恢复模型。

深度学习模型:CNN

CNN(卷积神经网络),适合用在图像识别和分类中,通过一系列的卷积层池化层全连接层来处理数据

在卷积层中,CNN使用一组可学习的滤波器(卷积核)来扫描输入的图像或信号。每个滤波器都能够检测输入中的特定特征,如边缘或颜色斑点。通过这种方式,CNN能够捕捉到图像中的局部特征,并保持这些特征的空间关系。

池化层(也称为下采样层)则用于降低特征的空间尺寸,从而减少参数数量和计算复杂度,同时使特征检测更加鲁棒。全连接层则将学习到的高级特征用于分类或其他任务。

💡CNN的结构和工作原理

一个典型的CNN包含以下几个主要部分:

  • 输入层:接收原始数据,如图像的像素值。

  • 卷积层:使用多个卷积核提取输入的特征。

  • 激活函数:如ReLU,用于引入非线性,使网络能够学习更复杂的特征。

  • 池化层:降低特征的空间维度,减少计算量。

  • 全连接层:将学习到的特征映射到最终的输出,如分类标签。

  • 输出层:输出网络的最终结果,如分类的概率分布。

CNN通过这些层的堆叠,能够从简单到复杂逐渐提取图像的特征。在训练过程中,CNN通过反向传播算法调整卷积核中的权重,以最小化预测结果和真实标签之间的差异。

训练MNIST

❓什么是MNIST

全称:(Modified National Institute of Standards and Technology),是机器学习和深度学习领域最经典的入门数据集,被称为深度学习的“Hello World”。它包含手写数字(0-9) 的灰度图像,广泛用于图像分类算法的训练与测试。

模型定义

首先定义模型的结构,这个结构在训练和推理过程中都是要保持一致的,其中包括模型组件的定义和模型流程定义,代码如下所示:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

上面的代码中,首先定义了模型的组件结构,其次forward是模型的流程结构。这是一个经典的“特征提取 + 分类器”双阶段结构,分为两个阶段。

第一阶段:特征提取(卷积神经网络),负责把像素变成“语义特征”。

  • 卷积层 1(Conv1):初级特征提取(边缘、线条)

    • self.conv1 = nn.Conv2d(1, 32, 3, 1)
    • in_channels=1 : 输入是灰度图(1个通道)
    • out_channels=32: 派出32个不同的“侦探”(滤波器)
    • kernel_size=3 : 每个侦探拿3x3的放大镜
    • stride=1 : 每次移动1个像素
    • 输出尺寸计算:(28 - 3 + 1) = 26 -> 输出 (32, 26, 26)
  • 卷积层 2:高级特征提取(部件、纹理)

    • self.conv2 = nn.Conv2d(32, 64, 3, 1)
    • in_channels=32: 接收上一层的32个特征图
    • out_channels=64: 增加到64个侦探,组合更复杂的模式
    • 输出尺寸计算:(26 - 3 + 1) = 24 -> 输出 (64, 24, 24)
  • Dropout 层 1:针对卷积特征的正则化

    • self.dropout1 = nn.Dropout2d(0.25)
    • 作用:随机“掐断”25%的通道(整张特征图),强迫模型不要把赌注押在某几个特定的特征组合上
    • 为什么用 Dropout2d:防止特征图之间产生共适应(Co-adaptation)
  • Dropout 层 2:针对全连接层的正则化

    • self.dropout2 = nn.Dropout2d(0.5)
    • 作用:随机“掐断”50%的神经元,因为全连接层参数最多,最容易过拟合

第二阶段:分类器(全连接网络),负责把特征变成具体的数字类别

  • 全连接层 1:特征映射与降维

    • self.fc1 = nn.Linear(9216, 128)
    • 输入维度 9216 的由来(关键计算)
    • 经过 Conv2(64, 24, 24) -> MaxPool(2, 2)
    • 输出尺寸:64个通道 * 12 * 12 = 9216
    • 输出维度 128:将高维特征压缩成128维的“数字指纹”
  • 全连接层 2:最终输出层

    • self.fc2 = nn.Linear(128, 10)
    • 输入 128:接收上层的指纹
    • 输出 10:对应 0-9 十个数字类别

而这些阶段落实到的流程执行中,需要用到一些额外的手段,如下所示:

  • relu,激活函数:引入非线性,筛选有效特征
  • max_pool2d,最大池化,模型不再关心像素的精确位置,只关心“有没有”
  • flatten,展平:把三维特征图拉直成一维向量,为了对接全连接层(只能吃一维数据)
  • log_softmax,转换为对数概率

forward函数清晰地展示了数据流(Flow):像素 → 边缘 → 部件 → 抽象特征 → 概率

❓ 我当时真实的困惑

  1. 为什么卷积核数量是 32和 64?不能是 30吗?
  2. 9216这个数字是怎么来的?
  3. forward 里,数据尺寸到底是怎么一步步变化的?

这几个是非常典型、也非常值得认真回答的问题。

🔑首先回答第一个问题,可以是 30,但几乎没人这么做。

原因有三个,都是工程现实,不是玄学:

1️⃣ 最重要的是GPU 硬件对齐,因为现代 GPU 的运算单元(CUDA Core / Tensor Core)是按 2 的幂次并行工作的,能刚好对齐显存带宽和线程块(block),如果用 30会出现 padding / waste,计算效率下降,训练变慢。

2️⃣ 经验性容量设计,在第一层提取简单特征(边、角、点),32个通道足够覆盖常见低级模式;第二层组合成复杂结构(弧线、圈、部件),需要更多表示能力,64是实践中验证过的稳定选择

3️⃣ 历史惯性(LeNet 传统),很多教程只是继承这种 翻倍增长​ 的设计,而不是重新发明

🔑其次是第二个问题,9216 不是玄学,是尺寸推导,是算出来的

✅ 已知前提

  • 输入图片:28 × 28
  • 卷积核:3 × 3
  • padding:0
  • stride:1

✅ 每一层尺寸变化

操作 输出尺寸
Input 原始图片 1 × 28 × 28
Conv1 3×3 卷积,无 padding 32 × 26 × 26
MaxPool 2×2 池化 32 × 13 × 13
Conv2 3×3 卷积 64 × 11 × 11
MaxPool 2×2 池化 64 × 5 × 5

⚠️ 注意:不同实现可能略有差异,但常见版本最后会再接一个卷积或 padding,使最终特征图为 64 × 12 × 12

那么最终,9216 的来源就是

64 个通道
× 12 × 12 每个通道的特征图
= 9216

模型训练

整体来说,训练过程就是把一个 batch 的数据送进模型 → 算损失 → 反向传播 → 更新参数 → 记录日志,这是深度学习里“学习”发生的唯一地方

📌 PyTorch 的数据是“流式”的,不是一次性全塞进去

def train(model, device, train_loader, optimizer, epoch, losses):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

        losses.append(loss.item())

        if batch_idx % 100 == 0:
            print(
                f'Epoch {epoch} '
                f'[{batch_idx * len(data)}/{len(train_loader.dataset)} '
                f'({100. * batch_idx / len(train_loader):.0f}%)]\t'
                f'Loss: {loss.item():.6f}'
            )

函数参数表含义如下所示:



参数
 


是什么
 
model 你的 CNN 网络
device CPU / GPU
train_loader 数据加载器(自动分批)
optimizer 优化器(Adam / SGD)
epoch 当前是第几轮
losses 用来记录 loss 的列表

首先要做的就是让模型切换到训练模式,使用model.train(),它告诉模型Dropout层开启随机失活,BatchNorm层使用当前batch 的均值/方差

如果有GPU,那么data, target = data.to(device), target.to(device)就是把数据搬到GPU

因为PyTorch 默认梯度是累加的,所以使用optimizer.zero_grad()清空梯度,来给新的学习腾地方

而模型在这里的作用就是做前向传播,就是前面所将的流程:

Conv1 → ReLU → Conv2 → ReLU → Pool → Dropout → FC → Softmax 

模型的输出output的形状是(batch_size, 10)

计算损失使用loss = F.nll_loss(output, target),它衡量的是模型猜的有多离谱,其中nll_loss是负对数似然,target是真实标签(0~9),我们肯定是想要最小化损失函数

反向传播(Backward Pass)是深度学习最核心的一行代码,可以自动计算∂Loss∂w,然后把
梯度存到每个参数的 .grad属性里,核心是把“总错误”按责任比例分摊回去,对比正向传播,过程如下:

Loss
 ↓
∂Loss/∂FC
 ↓
∂FC/∂Pool
 ↓
∂Pool/∂Conv2
 ↓
∂Conv2/∂Conv1
 ↓
∂Conv1/∂W

🔑有反向传播就可以自动学习,而不用手动调参

梯度告诉了方向,然后是优化器(Optimizer Step)optimizer.step()负责迈步子,读取.grad,根据优化规则更新参数,其中Adam 做的是:θ=θ−η⋅m√v+ϵ,这是“真正动手改模型”的地方

大致的训练流程就是这样,那么现在开始实际的训练吧。

def main():
    # ==========================================
    # 1. 设备选择与配置
    # ==========================================
    # 自动检测是否有可用的 GPU(CUDA)
    # 如果有则用 GPU 加速训练,否则回退到 CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # ==========================================
    # 2. 模型与优化器初始化
    # ==========================================
    # 实例化我们定义的 CNN 模型
    # .to(device) 将模型的所有参数和缓冲区移动到指定设备(GPU/CPU)
    model = Net().to(device)
    
    # 使用 Adam 优化器
    # model.parameters() 告诉优化器:需要更新哪些参数
    # lr=0.001 是学习率,控制每次参数更新的步长
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # ==========================================
    # 3. 数据预处理管道
    # ==========================================
    # transforms.Compose 将多个变换组合成一个流水线
    transform = transforms.Compose([
        # 将 PIL 图像或 numpy.ndarray 转换为 PyTorch Tensor
        # 同时将像素值从 [0, 255] 缩放到 [0.0, 1.0]
        transforms.ToTensor(),
        
        # 标准化处理(Normalization)
        # 使用 MNIST 数据集的官方均值和标准差
        # 公式: output = (input - mean) / std
        # 作用: 使数据分布均值为0,方差为1,加速模型收敛
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # ==========================================
    # 4. 数据集与数据加载器
    # ==========================================
    # 加载 MNIST 训练数据集
    dataset = datasets.MNIST(
        './data',          # 数据存储路径
        train=True,        # 使用训练集(而非测试集)
        download=True,     # 如果数据不存在,自动从网上下载
        transform=transform  # 应用上面定义的数据预处理
    )
    
    # DataLoader 负责批量加载数据,并提供打乱、并行读取等功能
    train_loader = DataLoader(
        dataset,           # 要加载的数据集
        batch_size=64,     # 每个批次包含 64 个样本
        shuffle=True       # 每个 epoch 都打乱数据顺序(防止模型记忆顺序)
    )

    # ==========================================
    # 5. 训练循环
    # ==========================================
    # 用于记录每个 batch 的损失值,以便后续可视化
    losses = []

    # 训练 10 个 epoch(完整遍历数据集 10 次)
    for epoch in range(1, 11):
        # 调用我们定义的 train 函数
        # 将模型、设备、数据加载器、优化器、当前轮次和损失列表传入
        train(model, device, train_loader, optimizer, epoch, losses)

    # ==========================================
    # 6. 保存训练好的模型
    # ==========================================
    # 保存模型的参数(state_dict)
    # 只保存参数而不保存整个模型是推荐做法,节省空间且灵活
    torch.save(model.state_dict(), "mnist_cnn.pt")

    # ==========================================
    # 7. 可视化训练过程
    # ==========================================
    # 创建一个 8x5 英寸大小的画布
    plt.figure(figsize=(8, 5))
    
    # 绘制损失曲线
    # losses 列表记录了每个 batch 的损失值
    plt.plot(losses, label="Training Loss")
    
    # 设置坐标轴标签
    plt.xlabel("Batch")      # X轴:批次索引
    plt.ylabel("Loss")       # Y轴:损失值
    
    # 设置图表标题
    plt.title("MNIST CNN Training Loss")
    
    # 显示网格线,便于观察数值
    plt.grid(True)
    
    # 显示图例(对应 label="Training Loss")
    plt.legend()
    
    # 自动调整子图参数,使之填充整个画布
    plt.tight_layout()
    
    # 显示图形窗口
    plt.show()

经过训练循环过程后,模型的犯错程度在不断的下降,如下图所示:

image

到这里,基本的训练过程就讲差不多了,那么一个epoch 里到底发生了什么?

Epoch 1
 ├─ Batch 1: 猜 → 算错 → 改参数
 ├─ Batch 2: 猜 → 算错 → 改参数
 ├─ ...
 └─ Batch N: 猜 → 算错 → 改参数

Epoch 2
 ├─ 猜得更准一点
 ├─ Loss 更小
 └─ ...

Epoch 10
 └─ 基本稳定

💡训练函数的本质,就是用损失函数告诉模型“你错在哪,再用反向传播告诉它“该怎么改”

模型测试

我使用了Forward Hook(前向钩子)方法,来达到可视化推理的中间过程,如下所示:

features = {}
def hook(name):
    def fn(_, __, out):
        features[name] = out.detach()
    return fn

# 注册钩子
model.conv1.register_forward_hook(get_features('conv1'))
model.conv2.register_forward_hook(get_features('conv2'))
model.dropout1.register_forward_hook(get_features('dropout1'))

进行测试时,加载之前训练好的模型权重,并开启model.eval(),同时需要注意模型定义要和训练时保持一致

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Net().to(device)
model.load_state_dict(torch.load("mnist_cnn.pt", map_location=device))
model.eval()

模型接受的输入需要先预处理一下,因为MNIST的图片有一定规范,比如尺寸、颜色等

# ============ 图片预处理 ===============
img_path = "my_digit.png"
img = Image.open(img_path).convert("L")

# ✅ 白底黑字 → 黑底白字(MNIST 风格)
img = ImageOps.invert(img)

# resize 到 28×28
img = img.resize((28, 28), Image.LANCZOS)

# 输入预处理
transform = T.Compose([
    T.ToTensor(),
    T.Normalize((0.1307,), (0.3081,))
])
# 获得模型输入
input_tensor = transform(img).unsqueeze(0).to(device)

最后就是模型推理代码了

# ==========================================
# 推理模式(禁用梯度计算)
# 作用:节省显存、加快计算、关闭 Dropout
# ==========================================
with torch.no_grad():
    
    # ==========================================
    # 前向传播:让模型看这张图片
    # output 形状: (1, 10)
    # 注意:这是 log_softmax 输出(对数概率)
    # ==========================================
    output = model(input_tensor)
    
    # ==========================================
    # 将对数概率转换回普通概率
    # 公式: prob = exp(log_prob)
    # 作用:让人类能直观理解(0~1 之间)
    # ==========================================
    prob = torch.exp(output)
    
    # ==========================================
    # 找出概率最大的类别(预测结果)
    # argmax(dim=1): 在类别维度上找最大值索引
    # item(): 把 Tensor 转成 Python 整数
    # 结果: 0~9 的数字
    # ==========================================
    pred = prob.argmax(dim=1).item()
    
    # ==========================================
    # 取出该类别的预测置信度
    # prob[0, pred]: 第 0 个样本、预测类别上的概率值
    # item(): 转成 Python 浮点数
    # 结果: 0.0~1.0 之间的置信度
    # ==========================================
    conf = prob[0, pred].item()

使用上面提到的钩子函数获取过程特征图,以Conv1举例:

conv1_feat = features['conv1'].squeeze().cpu()

plt.figure(figsize=(12, 6))
for i in range(32):
    plt.subplot(4, 8, i+1)
    plt.imshow(conv1_feat[i], cmap='gray')
    plt.axis('off')
plt.suptitle("Conv1: Edge Detectors")
plt.show()

下面就是推理的中间过程图片,有4个,分别是conv1,conv2,dropout1以及全连接层的输出,如下所示:

  1. Conv1 输出

32 张特征图,有的对竖边敏感,有的对横边敏感,有的对角点敏感

image

  1. Conv2 输出

不再是简单边缘,开始出现:圈、弧线、局部形状

image

  1. Pooling + Dropout输出

池化在“做减法”(精简信息),Dropout 在“做干扰”(增强韧性)

image

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐