PyTorch 卷积神经网络示例

本文详细介绍如何基于 PyTorch CPU 构建一个卷积神经网络 (CNN) 来对 CIFAR-10 图片进行分类。

介绍

• 构建一个经典的卷积神经网络 (CNN, convolution neural network)，对 CIFAR-10 数据集中的图像进行分类
• CNN 是一类神经网络，定义为用于检测数据中复杂特征的多层神经网络，它最常用于计算机视觉应用
• 本次示例 CNN 网络结构将包括以下 14 层：

Conv -> BatchNorm -> ReLU -> Conv -> BatchNorm -> ReLU -> MaxPool -> Conv -> BatchNorm -> ReLU -> Conv -> BatchNorm -> ReLU -> Linear.

• CIFAR-10 数据集是一个常用的计算机视觉数据集，包含 60000 张 32x32 彩色图片，分为 10 个类别，每个类别有 6000 张图片。其中 50000 张用于训练，10000 张用于测试。

实现

整个训练流程将包括以下几个步骤：

1. 准备环境：安装必要的库
2. 数据加载与预处理：下载 CIFAR-10 数据集，并进行必要的转换（如如将图片转换为 PyTorch 的 Tensor 格式，并进行归一化）
3. 定义神经网络模型：构建一个基于卷积层的神经网络
   1. 卷积神经网络：通常由卷积层 (Conv2d)、激活函数 (ReLU)、池化层 (MaxPool2d) 和全连接层 (Linear) 组成
4. 定义损失函数与优化器：选择合适的损失函数和优化算法
   1. 损失函数，如 Adam、交叉熵损失 (CrossEntropyLoss) 作为损失函数
5. 训练模型：在训练集上迭代训练模型
6. 评估模型：在测试集上评估模型的性能

项目目录结构：

$ tree .
.
├── cnn-demo.pth
├── data # 数据集可以自动下载
│   └── cifar-10-python.tar.gz
└── PyTorchTraining.py

准备环境

确保你已经安装了 PyTorch。如果没有，可以使用 pip 安装：

pip install torch torchvision matplotlib numpy===1.26.4

下载数据集

在 https://www.cs.toronto.edu/~kriz/cifar.html 下载 cifar-10-python.tar.gz，放到 data 目录下

源码

• PyTorchTraining.py

source repo / download raw

训练模型

将迭代多个 epoch，在每个 epoch 中遍历训练数据，执行前向传播、计算损失、反向传播和参数更新。

$ python3 PyTorchTraining.py
Files already downloaded and verified
The number of images in a training set is:  50000
Files already downloaded and verified
The number of images in a test set is:  10000
The number of batches per epoch is:  5000
神经网络模型定义完成。
损失函数和优化器定义完成。
开始训练模型，共 5 个 epoch...
The model will be running on cpu device
[1/5,  1000] loss: 1.764
[1/5,  2000] loss: 1.472
[1/5,  3000] loss: 1.303
[1/5,  4000] loss: 1.168
[1/5,  5000] loss: 1.094
For epoch 1 the test accuracy over the whole test set is 64 %

[2/5,  1000] loss: 1.027
[2/5,  2000] loss: 1.001
[2/5,  3000] loss: 0.981
[2/5,  4000] loss: 0.983
[2/5,  5000] loss: 0.941
For epoch 2 the test accuracy over the whole test set is 67 %

[3/5,  1000] loss: 0.868
[3/5,  2000] loss: 0.865
[3/5,  3000] loss: 0.865
[3/5,  4000] loss: 0.868
[3/5,  5000] loss: 0.859
For epoch 3 the test accuracy over the whole test set is 69 %

[4/5,  1000] loss: 0.802
[4/5,  2000] loss: 0.785
[4/5,  3000] loss: 0.785
[4/5,  4000] loss: 0.792
[4/5,  5000] loss: 0.777
For epoch 4 the test accuracy over the whole test set is 70 %

[5/5,  1000] loss: 0.721
[5/5,  2000] loss: 0.719
[5/5,  3000] loss: 0.739
[5/5,  4000] loss: 0.744
[5/5,  5000] loss: 0.740
For epoch 5 the test accuracy over the whole test set is 71 %

Finished Training
模型结构：
  conv1: Conv2d(3, 12, kernel_size=(5, 5), stride=(1, 1), padding=(1, 1))
  bn1: BatchNorm2d(12, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  conv2: Conv2d(12, 12, kernel_size=(5, 5), stride=(1, 1), padding=(1, 1))
  bn2: BatchNorm2d(12, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  pool: MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  conv4: Conv2d(12, 24, kernel_size=(5, 5), stride=(1, 1), padding=(1, 1))
  bn4: BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  conv5: Conv2d(24, 24, kernel_size=(5, 5), stride=(1, 1), padding=(1, 1))
  bn5: BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  fc1: Linear(in_features=2400, out_features=10, bias=True)
Real labels:    cat  ship  ship plane  frog  frog   car  frog   cat   car
Predicted:    dog   car plane plane  deer   dog   car  deer   dog   car

• 貌似预测的一点也不准，有时间试试这个
• 该模型上传到 https://huggingface.co/xiexianbin/cnn-demo

结束

通过上述代码和详细注释，你应该能够理解如何基于 PyTorch CPU 构建和训练一个简单的卷积神经网络来完成 CIFAR-10 图片分类任务。你可以尝试修改网络结构、超参数（如学习率、epoch 数量、批次大小）来进一步优化模型性能。