60分钟实战速通PyTorch：从零构建图像分类器

Mu Tian

310人浏览 · 2026-07-04 10:24:16

Mu Tian · 2026-07-04 10:24:16 发布

1. 环境准备与工具安装

想用PyTorch玩转图像分类？先别急着写代码，咱们得把工具配齐了。就像做饭前要备好锅碗瓢盆一样，搞深度学习也得先准备好开发环境。我推荐直接用Anaconda来管理Python环境，省得折腾各种依赖问题。

装PyTorch其实特别简单，打开终端输入下面这行命令（如果你用GPU加速的话记得选CUDA版本）：

conda install pytorch torchvision torchaudio -c pytorch

验证安装是否成功可以跑个简单测试：

import torch
print(torch.__version__)  # 应该输出1.x.x这样的版本号
print(torch.cuda.is_available())  # 检查GPU是否可用

这里有个小坑我踩过——如果你的显卡是AMD的，可能得用ROCm版本的PyTorch。不过新手建议先用CPU版本来练手，等熟悉了再折腾GPU加速。

数据集我们选用经典的CIFAR-10，这个数据集已经内置在torchvision里了，包含6万张32x32的彩色图片，分为10个类别（飞机、汽车、鸟之类的）。为啥选它？因为MNIST太简单了，ImageNet又太大，CIFAR-10正好适合练手。

2. 数据加载与预处理

PyTorch的数据处理有个标准流程：先用Dataset定义数据，再用DataLoader批量加载。咱们先看看怎么把图片数据整理成模型能吃的格式：

import torchvision
import torchvision.transforms as transforms

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),  # 转成Tensor格式
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

# 下载并加载训练集
trainset = torchvision.datasets.CIFAR10(
    root='./data', 
    train=True,
    download=True, 
    transform=transform
)
trainloader = torch.utils.data.DataLoader(
    trainset, 
    batch_size=32,
    shuffle=True
)

# 同样的方法加载测试集
testset = torchvision.datasets.CIFAR10(
    root='./data',
    train=False,
    download=True,
    transform=transform
)
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=32,
    shuffle=False
)

这里有几个实用技巧：

shuffle=True 会让每轮训练的数据顺序都打乱，防止模型记住顺序
batch_size 一般设为32或64，太小了训练慢，太大了可能内存不够
归一化用(0.5, 0.5, 0.5)是因为CIFAR-10图片的RGB通道均值大约都在0.5左右

想看下数据长啥样？可以用matplotlib可视化：

import matplotlib.pyplot as plt
import numpy as np

# 显示图片的函数
def imshow(img):
    img = img / 2 + 0.5  # 反归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 获取一批训练数据
dataiter = iter(trainloader)
images, labels = next(dataiter)

# 显示图片
imshow(torchvision.utils.make_grid(images[:4]))
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

3. 构建卷积神经网络

现在来搭建我们的CNN模型。别被"卷积神经网络"吓到，其实就像搭积木一样简单。PyTorch提供了 nn.Module 这个基类，我们只需要继承它并实现两个方法：

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        # 卷积层
        self.conv1 = nn.Conv2d(3, 6, 5)  # 输入通道3，输出通道6，卷积核5x5
        self.pool = nn.MaxPool2d(2, 2)  # 池化层
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 全连接层
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)  # 10类输出

    def forward(self, x):
        # 前向传播
        x = self.pool(F.relu(self.conv1(x)))  # 卷积→激活→池化
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)  # 展平
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

这个网络结构是这样的：

输入图片(3,32,32) → 卷积1 → (6,28,28) → 池化 → (6,14,14)
→ 卷积2 → (16,10,10) → 池化 → (16,5,5)
→ 展平 → 400维向量
→ 全连接层1 → 120维 → 全连接层2 → 84维 → 输出10维

为什么要用ReLU激活函数？因为它能有效缓解梯度消失问题，而且计算简单。池化层的作用是降维，保留主要特征的同时减少计算量。

4. 训练模型

模型搭好了，现在进入最激动人心的训练环节！我们需要定义损失函数和优化器：

import torch.optim as optim

criterion = nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  # 带动量的SGD

训练循环的代码长这样：

for epoch in range(10):  # 跑10轮
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入数据
        inputs, labels = data
        
        # 梯度清零
        optimizer.zero_grad()
        
        # 前向传播+反向传播+优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 打印统计信息
        running_loss += loss.item()
        if i % 500 == 499:  # 每500个batch打印一次
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 500))
            running_loss = 0.0

print('训练结束')

这里有几个关键点：

zero_grad() 必须调用，否则梯度会累积
loss.backward() 会自动计算所有参数的梯度
optimizer.step() 根据梯度更新参数
学习率lr是个重要超参数，太大容易震荡，太小收敛慢

训练过程中你会看到loss逐渐下降，这说明模型在学习。如果loss不降反升，可能是学习率设太大了。

5. 模型测试与保存

训练完的模型需要评估性能。我们用测试集来检查模型的泛化能力：

correct = 0
total = 0
with torch.no_grad():  # 测试时不需要计算梯度
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('测试集准确率: %d %%' % (100 * correct / total))

保存模型也很简单：

PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

加载模型时先实例化网络，再加载参数：

net = Net()  # 必须先创建相同结构的网络
net.load_state_dict(torch.load(PATH))

如果想用GPU加速，只需要把模型和数据都转到CUDA上：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
inputs = inputs.to(device)
labels = labels.to(device)

6. 常见问题与调优技巧

新手常遇到几个坑：

梯度爆炸 ：如果出现NaN，尝试减小学习率或用梯度裁剪
过拟合 ：可以添加Dropout层或L2正则化
训练慢 ：适当增大batch size或用GPU加速

改进模型性能的几个方法：

增加网络深度（但别太深，CIFAR-10用3-5层卷积就够了）
使用更先进的优化器如Adam
添加Batch Normalization层
尝试数据增强（随机翻转、裁剪等）

# 改进版网络示例
class ImprovedNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x