登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 业界 ›从零开始：基于 PyTorch 的图像分类模型

从零开始：基于 PyTorch 的图像分类模型

宁觅波 3 天前
摘要：本文详细记录了使用 PyTorch 从零搭建一个图像分类模型的过程，涵盖卷积神经网络（CNN）、数据预处理、模型设计、训练调试与优化。通过对 CIFAR-10 数据集的处理实践，结合经典文献和 2025 年最新研究趋势，深入探讨了技术细节，并辅以完整实践源码的过程和结论。我选择用 PyTorch 搭建图像分类模型，既源于对深度学习的兴趣，也因为它在 2025 年的技术社区中热度不减。通过这次实践，我希望掌握 CNN 的核心原理，同时记录过程，为其他初学者提供参考。
关键字：PyTorch、图像分类、CNN、深度学习、模型优化、CIFAR-10、调试经验、前沿趋势
引言

背景介绍

图像分类是计算机视觉的核心任务，广泛应用于自动驾驶、医疗影像分析和人脸识别等领域。
深度学习，特别是卷积神经网络（CNN），极大推动了这一技术的发展。2012年，AlexNet在ImageNet挑战赛中大幅降低分类错误率，标志着深度学习时代的开端。
此后，模型架构从 ResNet 进化到 VisionTransformer（ViT），性能不断提升。到2025年，硬件算力增强和PyTorch 2.x 等框架的优化（如torch.compile()），让图像分类任务更加高效。在实践中，以CIFAR-10等数据集为例，一个简单CNN在现代GPU上几分钟即可训练完成，展现了技术的进步。
本文目标与结构


  • 目标：搭建并优化一个 CNN，完成 CIFAR-10 分类任务。
  • 结构：理论讲解 → 源码实现 → 调试分析 → 优化与展望。
卷积神经网络（CNN）

为何需要 CNN？

传统人工神经网络（ANN）在图像处理上参数量巨大，计算成本高。CNN 通过局部感受野、权重共享和层级特征提取（本文只入门，这三个都不涉及），减少计算量并提升模型能力。
CNN 发展始于 LeNet-5（1998），后续有 AlexNet（2012）、VGG（2014）、ResNet（2015），极大推动计算机视觉进步。
卷积神经网络（Convolutional Neural Network, CNN）是深度学习中处理图像任务的基石，其核心在于通过卷积操作提取空间特征、优化参数，并逐步构建高级语义表示。
CNN 的核心组成

1、卷积层（Convolutional Layer）  ：特征提取的“智能滤镜”

CNN 的核心是卷积操作，通过一个小的滑动窗口（卷积核, kernel）在输入图像上扫描，计算局部区域的加权和，提取边缘、纹理等低级特征。
卷积包括填充、步幅等参数，还包括和全连接层一样的偏置参数，卷积运算相当于图像处理中的“滤波器运算”。
卷积运算会先对输入数据进行初始化准备。数据除了高、长方向之外，还需要处理通道方向，形成三维数据(Channel, Height, Width)。再结合批处理参数，按照(batch_num, channel, height, width)顺序保存数据，形成四维数据。
卷积运算通过填充、步幅对数据进行再加工。填充即向输入数据的周围填入0，主要是为了调整输出的大小。避免因每次卷积运算缩小的数据空间最终导致卷积运算不可用的情况。步幅是卷积运算每一次运算的数据位置间隔，通过步幅计算出输出数据矩阵大小。
卷积运算通过卷积核完成数据运算。卷积核即滤波器，是一种“滤镜”，通过它来完成指定运算目标（识别、聚焦等）。
卷积运算通过偏置函数完成数据输出。偏置是在数据输出前针对数据的统一处理，以适配特定场景（整体调亮等）。
简单来说，卷积就像用一个“放大镜”在图像上滑动，每次聚焦一小块，找出关键线索。
卷积操作类似于“盲人摸象”的工作过程，举个例子：假设你看一张猫的照片，第一次卷积可能找到耳朵的边缘，第二次找到毛发的纹理，最终拼凑出“猫”的完整轮廓。这种层层递进的过程，正是卷积的魅力。
①输入数据和卷积核

假设输入图像是一个 4x4 的灰度矩阵（单通道），卷积核为 3x3，步幅为 1，无填充。

  • 输入矩阵

\[I = \begin{bmatrix}1 & 2 & 3 & 0 \\0 & 1 & 2 & 3 \\3 & 0 & 1 & 2 \\2 & 3 & 0 & 1\end{bmatrix}\]

  • 卷积核（滤波器）

\[K = \begin{bmatrix}2 & 0 & 1 \\0 & 1 & 2 \\1 & 0 & 2 \end{bmatrix}\]
②计算过程（可忽略）

卷积运算通过滑动卷积核，在输入矩阵上逐块计算点积。
输出特征图大小为2x2，I[高]-K[高]+步幅  *  I[宽]-K[宽]+步幅 ：

\[(4-3+1) \times (4-3+1) = 2 \times 2\]

  • 公式：对于输入I和卷积核K，输出O[i,j]为：

\[O[i,j]= m∑​  n∑​ I[i+m,j+n]⋅K[m,n]\]
③手动计算


  • 左上角 2+0+3+0+1+4+3+0+2=15
  • 右上角 4+0+0+0+2+6+0+0+4=16
  • 左下角 0+0+2+0+0+2+2+0+0=6
  • 右下角 2+0+3+0+1+4+3+0+2=15
输出特征图为：

\[O = \begin{bmatrix}15 & 16 \\6 & 15\end{bmatrix}\]
2、池化层（Pooling Layer）  ：降维与聚焦关键信息

池化是缩小高、长方向上的空间的运算。用于压缩特征图，减少计算量，同时保留重要信息。
LeCun 的LeNet-5引入了平均池化（Average Pooling）：计算目标区域的平均值。
AlexNet 则普及了最大池化（Max Pooling）：计算目标区域的最大值。
想象你在看一幅画，池化就像眯起眼睛，只关注最亮的亮点（最大值），忽略细枝末节。这样既降低了分辨率，又让网络更关注显著特征，比如猫脸上的胡须而非背景噪声。
池化的现代应用更灵活，也体现了CNN的进化，例如全局平均池化（Global Average Pooling）常用于替代全连接层，提升模型泛化能力。
最大池化示例（2×2 池化，步幅 2）：

将4x4的输入数据按照步幅2计算，输出为2x2，分为4个区域：

  • 左上：[1,2,5,6]，最大值6
  • 右上：[2,4,8,7]，最大值8
  • 左下：[9,10,13,15]，最大值15
  • 右下：[12,11,14,16]，最大值16
如下图所示：

\[\begin{bmatrix}1 & 3 & 2 & 4 \\5 & 6 & 8 & 7 \\9 & 10 & 12 & 11 \\13 & 15 & 14 & 16\end{bmatrix}\Rightarrow\begin{bmatrix}6 & 8 \\15 & 16\end{bmatrix}\]
以上为Max池化，如果是Average池化，则改为平均值即可。不赘述。
3、激活函数（Activation Function）  ：赋予非线性表达力

卷积和池化后的特征需要通过激活函数（Activation Function）引入非线性，否则网络只能学到线性变换。
ReLU 是 CNN 最常用的非线性变换：

\[f(x) = \max(0, x)\]
它计算高效，并避免梯度消失问题。
ReLU 的作用像个“开关”：负值关掉，正值保留，让网络学会更复杂的模式，比如区分猫和狗的不同轮廓。
现代还有 Leaky ReLU 和 Swish 等变种，但 ReLU 仍是理解 CNN 非线性的起点。举例来说，ReLU 就像在筛选线索时，只保留“有用的证据”，丢弃无关信息。
4、全连接层（Fully Connected Layer, FC）  ：特征整合与分类

在经过卷积和池化后，CNN 将特征展平（Flatten）并输入全连接层，用于分类或回归任务。
全连接层就像大脑的“决策中心”，综合所有线索，给出最终判断：这张图是“猫”还是“狗”？
不过，现代趋势逐渐减少全连接层依赖，例如 ResNet（2015）用全局池化简化输出，减轻过拟合风险。
CNN 在计算机视觉中的应用


  • 图像分类（AlexNet、ResNet）
  • 目标检测（YOLO、Faster R-CNN）
  • 语义分割（U-Net、DeepLab）
  • 医学影像分析（CT 诊断）
实践环境与准备

工具与依赖


  • Python 3.10、PyTorch 2.1、torchvision 0.16。
  • 硬件：NVIDIA RTX 3060 GPU，16GB RAM。
Anaconda 环境配置

1. 安装 Anaconda

使用 Homebrew 安装 Anaconda：
  1. brew install --cask anaconda
复制代码
2. 初始化 Conda

将 conda 添加到 PATH：
  1. echo 'export PATH="/opt/homebrew/anaconda3/bin:$PATH"' >> ~/.zshrc
复制代码
初始化 conda 以在 zsh 中使用：
  1. conda init zsh
复制代码
重新加载 shell 配置：
  1. source ~/.zshrc
复制代码
3. 创建和配置项目环境

创建新的 conda 环境：
  1. conda create -n simpletorch python=3.9 -y
复制代码
激活环境：
  1. conda activate simpletorch
复制代码
安装项目依赖：
  1. pip install -r requirements.txt
复制代码
4. 验证环境

检查 Python 版本和已安装的包：
  1. python --version
  2. pip list | grep -E "torch|numpy|matplotlib"
复制代码
5. 常用命令


  • 激活环境：conda activate simpletorch
  • 退出环境：conda deactivate
  • 查看所有环境：conda env list
  • 删除环境：conda env remove -n simpletorch
6. 环境信息


  • 环境名称：simpletorch
  • Python 版本：3.9.21
  • 主要包版本：

    • PyTorch 2.6.0
    • torchvision 0.21.0
    • NumPy 2.0.2
    • Matplotlib 3.9.4

数据集选择与加载

CIFAR-10 数据集概述


  • 数据集大小：50,000 张图片
  • 图片大小：3x32x32（RGB图像）
  • 类别数：10个（飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车）
数据文件内容

data目录下的文件是 CIFAR-10 数据集文件，它们是在运行 main.py 或 show_dataset.py 时通过 torchvision.datasets.CIFAR10 自动下载的。具体来说，这些文件来自：

  • 官方下载地址：https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
  • 当代码执行到这一行时：
  1. trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
复制代码

  • root='./data' 指定下载到当前目录下的 data 文件夹
  • download=True 表示如果数据不存在就自动下载
  • 下载的文件会被自动解压到 data/cifar-10-batches-py/ 目录下
下载的文件包括：

  • data_batch_1 到 data_batch_5：训练数据
  • batches.meta：类别信息
  • test_batch：测试数据（虽然我们目前没有使用）
这些文件是数据集的一部分，通常不需要手动管理，PyTorch 会自动处理下载和解压过程。如果您想重新下载数据集，可以：

  • 删除 data 目录
  • 重新运行程序，它会自动重新下载
数据预处理
  1. transform = transforms.Compose([
  2.     transforms.ToTensor(),
  3.     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  4. ])
复制代码
数据集可视化

运行以下命令查看数据集样本和分布：
  1. python show_dataset.py
复制代码
这将生成两个可视化文件：

  • cifar10_samples.png：数据集样本图片
  • class_distribution.png：类别分布
源码分析
  1. import torch
  2. from torchvision import datasets, transforms
  3. import matplotlib.pyplot as plt
  5. def show_dataset():
  6.     # 数据预处理
  7.     transform = transforms.Compose([
  8.         transforms.ToTensor(),
  9.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  10.     ])
  12.     # 加载训练集
  13.     print("Loading CIFAR-10 dataset...")
  14.     trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
  15.    
  16.     # 定义类别
  17.     classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
  18.    
  19.     # 显示数据集信息
  20.     print(f"\nDataset size: {len(trainset)} images")
  21.     print(f"Image size: {trainset[0][0].shape}")
  22.     print(f"Number of classes: {len(classes)}")
  23.    
  24.     # 显示样本图片
  25.     plt.figure(figsize=(15, 5))
  26.     for i in range(5):
  27.         img, label = trainset[i]
  28.         img = img / 2 + 0.5
  29.         img = img.numpy()
  30.         plt.subplot(1, 5, i + 1)
  31.         plt.imshow(img.transpose(1, 2, 0))
  32.         plt.title(f'Class: {classes[label]}')
  33.         plt.axis('off')
  34.     plt.tight_layout()
  35.     plt.savefig('cifar10_samples.png')
  36.    
  37.     # 显示每个类别的样本数量
  38.     class_counts = torch.zeros(10)
  39.     for _, label in trainset:
  40.         class_counts[label] += 1
  41.    
  42.     plt.figure(figsize=(10, 5))
  43.     plt.bar(classes, class_counts)
  44.     plt.title('Number of samples per class')
  45.     plt.xticks(rotation=45)
  46.     plt.tight_layout()
  47.     plt.savefig('class_distribution.png')
  48.    
  49.     print("\nVisualization files saved:")
  50.     print("- cifar10_samples.png")
  51.     print("- class_distribution.png")
  53. if __name__ == '__main__':
  54.     show_dataset()
复制代码
运行结果
  1. Loading CIFAR-10 dataset...
  3. Dataset size: 50000 images
  4. Image size: torch.Size([3, 32, 32])
  5. Number of classes: 10
  7. Visualization files saved:
  8. - cifar10_samples.png
  9. - class_distribution.png
复制代码
模型设计与源码分析

网络结构设计


  • 结构：2 层卷积 + 2层池化 + 1 层全连接。
  • 源码：
    1. import torch
    2. import torch.nn as nn
    4. class SimpleNet(nn.Module):
    5.     def __init__(self):
    6.         super(SimpleNet, self).__init__()
    7.         # 第一个卷积层
    8.         # 输入: 3通道(RGB图像) -> 输出: 16个特征图
    9.         # 例如: 一张32x32的彩色图片(3x32x32) -> 16个32x32的特征图
    10.         self.conv1 = nn.Conv2d(3, 16, 3, padding=1)  # 448参数
    11.         
    12.         # 第二个卷积层
    13.         # 输入: 16个特征图 -> 输出: 32个特征图
    14.         # 例如: 16个16x16的特征图 -> 32个16x16的特征图
    15.         self.conv2 = nn.Conv2d(16, 32, 3, padding=1)  # 4,640参数
    16.         
    17.         # 全连接层
    18.         # 输入: 32 * 8 * 8 = 2048个特征 -> 输出: 10个类别
    19.         # 例如: 32个8x8的特征图展平后 -> 10个数字(0-9)的概率
    20.         self.fc1 = nn.Linear(32 * 8 * 8, 10)  # 20,490参数
    22.     def forward(self, x):
    23.         # 输入x的形状: [batch_size, 3, 32, 32]
    24.         # 例如: [32, 3, 32, 32] 表示32张32x32的RGB图片
    25.         
    26.         # 第一个卷积层 + ReLU激活
    27.         # 输出形状: [batch_size, 16, 32, 32]
    28.         # 例如: [32, 16, 32, 32] 表示32张图片，每张有16个32x32的特征图
    29.         x = torch.relu(self.conv1(x))
    30.         
    31.         # 最大池化层，将特征图尺寸减半
    32.         # 输出形状: [batch_size, 16, 16, 16]
    33.         # 例如: [32, 16, 16, 16] 表示32张图片，每张有16个16x16的特征图
    34.         x = torch.max_pool2d(x, 2)
    35.         
    36.         # 第二个卷积层 + ReLU激活
    37.         # 输出形状: [batch_size, 32, 16, 16]
    38.         # 例如: [32, 32, 16, 16] 表示32张图片，每张有32个16x16的特征图
    39.         x = torch.relu(self.conv2(x))
    40.         
    41.         # 最大池化层，再次将特征图尺寸减半
    42.         # 输出形状: [batch_size, 32, 8, 8]
    43.         # 例如: [32, 32, 8, 8] 表示32张图片，每张有32个8x8的特征图
    44.         x = torch.max_pool2d(x, 2)
    45.         
    46.         # 将特征图展平成一维向量
    47.         # 输出形状: [batch_size, 32 * 8 * 8]
    48.         # 例如: [32, 2048] 表示32张图片，每张图片的特征被展平成2048维向量
    49.         x = x.view(x.size(0), -1)
    50.         
    51.         # 全连接层，得到最终的类别预测
    52.         # 输出形状: [batch_size, 10]
    53.         # 例如: [32, 10] 表示32张图片，每张图片对应10个类别的预测概率
    54.         x = self.fc1(x)
    55.         return x
    复制代码
参数与设计分析


  • 参数总体情况：这是一个"轻特征提取，重分类"的模型，总参数量25,578

    • 全连接层参数占大多数（约80%）
    • 卷积层参数相对较少（约20%）
    • 总参数量适中，适合入门学习

  • 卷积核：3x3卷积核

    • 计算效率高（9个参数）
    • 感受野适中（能捕捉局部特征）
    • 是CNN中的标准选择

  • Padding填充设计：padding=1，输入32x32，输出32x32

    • 保持特征图尺寸不变
    • 避免边缘信息丢失
    • 便于网络设计

3x3卷积核 + padding=1）是CNN中的经典配置，既保证了特征提取效果，又维持了计算效率。
整体设计分析，最关键的三点如下：

  • 使用3x3卷积核配合padding=1，在保持特征图尺寸的同时高效提取局部特征。
  • 通过两层卷积（3→16→32通道）和两次池化（32x32→16x16→8x8），实现了从基本特征到复杂特征的渐进提取。
  • 最后用全连接层（2048→10）直接分类，总参数量25,578，结构简单但有效。
训练与调试

训练流程


  • 源码：
    1.                 # 训练一个epoch
    2.     print("Training for 1 epoch...")
    3.     model.train()  # 设置为训练模式（启用dropout等训练特定层）
    4.    
    5.     # 遍历数据加载器
    6.     for i, (inputs, labels) in enumerate(trainloader):
    7.         # 将数据移到指定设备（GPU/CPU）
    8.         inputs, labels = inputs.to(device), labels.to(device)
    9.         
    10.         # 前向传播
    11.         optimizer.zero_grad()  # 清空梯度
    12.         outputs = model(inputs)  # 模型前向传播
    13.         loss = criterion(outputs, labels)  # 计算损失
    14.         loss.backward()  # 反向传播，计算梯度
    15.         optimizer.step()  # 更新模型参数
    17.         # 每100个batch打印一次损失
    18.         if (i + 1) % 100 == 0:
    19.             print(f'Batch [{i + 1}], Loss: {loss.item():.4f}')
    21.     print("Training finished!")
    复制代码
调试经验


  • 数据先行：检查数据加载和预处理，确保输入数据的形状、范围和分布符合预期。
  • 监控为王：密切关注损失值变化趋势，通过定期打印损失、学习率和梯度信息来诊断训练状态。
  • 调参有度：根据训练效果合理调整超参数，如batch_size、学习率等，避免过拟合或欠拟合。
测试执行与结果分析

测试执行
  1. python main.py
复制代码
结果与分析
  1. Using device: cpu
  2. Loading CIFAR-10 dataset...
  3. Training for 1 epoch...
  4. Batch [100], Loss: 2.2419
  5. Batch [200], Loss: 2.2013
  6. Batch [300], Loss: 1.8651
  7. Batch [400], Loss: 1.9359
  8. Batch [500], Loss: 1.9718
  9. Batch [600], Loss: 1.9448
  10. Batch [700], Loss: 1.7974
  11. Batch [800], Loss: 1.6378
  12. Batch [900], Loss: 1.5137
  13. Batch [1000], Loss: 1.5045
  14. Batch [1100], Loss: 1.8072
  15. Batch [1200], Loss: 1.9754
  16. Batch [1300], Loss: 1.8177
  17. Batch [1400], Loss: 1.7377
  18. Batch [1500], Loss: 1.8140
  19. Training finished!
复制代码

  • 设备使用：程序在 CPU 上运行
  • 数据集：成功加载了 CIFAR-10 数据集
  • 训练过程：

    • 完成了 1 个 epoch 的训练
    • 每 100 个 batch 打印一次 loss
    • Loss 值从初始的 2.24 逐渐下降到约 1.81，说明模型在学习
    • 整个训练过程顺利完成

Loss（损失）是衡量模型预测结果与真实值之间差距的指标，就像考试打分一样 - 比如模型预测一张图片是猫的概率为60%，而实际上确实是猫（100%），这个40%的差距就反映在loss值上，loss越小代表模型预测越准确。
优化与前沿探索

架构优势


  • 轻量级设计：总参数量约25K，适合快速部署和迭代
  • 结构清晰：采用经典的CNN+池化层组合，便于理解和优化
  • 模块化实现：代码组织合理，便于扩展
与前沿趋势的差距


  • 缺少注意力机制（如Transformer结构）
  • 没有使用残差连接（ResNet特性）
  • 缺乏正则化策略（如Dropout）
可优化方向


  • 添加BatchNorm提高训练稳定性
  • 引入现代激活函数（如GELU、Swish）
  • 实现学习率调度策略
总结与反思

收获与不足


  • 掌握 CNN 与 PyTorch 基本流程。
  • 不足：模型深度有限，准确率待提升。
下一步计划


  • 添加验证集评估和模型保存机制
  • 实现训练过程可视化
从玩具到工具，从黑盒到透明
项目源码

参考文献


  • LeCun, Y., Bottou, L., Bengio, Y., & Haffner, P. (1998). Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11), 2278-2324. DOI:10.1109/5.726791
  • Krizhevsky, A., Sutskever, I., & Hinton, G. E. (2012). ImageNet classification with deep convolutional neural networks. Advances in Neural Information Processing Systems, 25, 1097-1105. DOI:10.1145/3065386
  • Simonyan, K., & Zisserman, A. (2014). Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556. arXiv:1409.1556
扫码关注公众号
1.jpeg

更多文章请转到一面千人的博客园


来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

使用道具 举报
相关推荐
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
宁觅波
3 天前

0

粉丝关注

12

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫