惚兮1. 项目概述
在计算机视觉领域，图像分类是最基础也最经典的任务之一。作为一名长期使用PyTorch框架的开发者，我发现很多初学者在入门图像分类时，往往被数据预处理、模型构建和训练调参这些环节卡住。本文将用最直白的语言，带大家完整走一遍PyTorch图像分类的全流程，从数据增强到CNN模型训练，每个步骤都会配上可运行的代码示例。
这个教程特别适合：
- 刚接触PyTorch和深度学习的小白
- 想系统学习图像分类pipeline的转行者
- 需要快速实现一个分类demo的开发者
我们将使用经典的CIFAR-10数据集，它包含10个类别的6万张32x32小图片，非常适合教学演示。整个流程可以在普通笔记本电脑上运行，不需要高端GPU。
需要模型API调用？ 免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 环境准备与数据加载
2.1 PyTorch环境配置
推荐使用Anaconda创建虚拟环境：
bash复制conda create -n pytorch_tutorial python=3.8
conda activate pytorch_tutorial
conda install pytorch torchvision torchaudio -c pytorch
验证安装是否成功：
python复制import torch
print(torch.__version__) # 应该输出2.x.x
print(torch.cuda.is_available()) # 检查GPU是否可用
注意：如果使用GPU训练，请确保安装了对应版本的CUDA驱动。PyTorch官网提供了详细的版本匹配指南。
2.2 数据加载与可视化
PyTorch的torchvision已经内置了CIFAR-10数据集：
python复制import torchvision
import torchvision.transforms as transforms
transform = transforms.Compose([
transforms.ToTensor(),
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=4, shuffle=True, num_workers=2)
classes = ('plane', 'car', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck')
可视化一批训练数据：
python复制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))
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))
3. 数据增强实战技巧
3.1 常用数据增强方法
数据增强能有效防止过拟合，特别是在小数据集上。我们扩展之前的transform：
python复制train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.RandomRotation(10), # 随机旋转±10度
transforms.ColorJitter( # 颜色抖动
brightness=0.1, contrast=0.1, saturation=0.1),
transforms.RandomAffine( # 随机仿射变换
degrees=0, translate=(0.1, 0.1)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
实操心得：数据增强的顺序很重要。几何变换（翻转、旋转）应该放在颜色变换之前，且ToTensor()和Normalize()通常放在最后。
3.2 自定义数据增强
PyTorch支持自定义增强策略。例如实现Cutout增强：
python复制class Cutout(object):
def __init__(self, length):
self.length = length
def __call__(self, img):
h, w = img.size(1), img.size(2)
mask = np.ones((h, w), np.float32)
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - self.length // 2, 0, h)
y2 = np.clip(y + self.length // 2, 0, h)
x1 = np.clip(x - self.length // 2, 0, w)
x2 = np.clip(x + self.length // 2, 0, w)
mask[y1:y2, x1:x2] = 0.
mask = torch.from_numpy(mask)
mask = mask.expand_as(img)
img *= mask
return img
使用时添加到transform中：
python复制train_transform.transforms.insert(-2, Cutout(8)) # 在ToTensor前插入
4. CNN模型构建与训练
4.1 基础CNN模型设计
我们实现一个简化版ResNet：
python复制import torch.nn as nn
import torch.nn.functional as F
class BasicBlock(nn.Module):
def __init__(self, in_planes, planes, stride=1):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(
in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(
planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes)
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, num_blocks, num_classes=10):
super(ResNet, self).__init__()
self.in_planes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1]*(num_blocks-1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes
return nn.Sequential(*layers)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
def ResNet18():
return ResNet(BasicBlock, [2,2,2,2])
4.2 训练流程实现
完整的训练循环：
python复制import torch.optim as optim
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = ResNet18().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
for epoch in range(10): # 实际可以训练更多轮
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data[0].to(device), data[1].to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 2000 == 1999: # 每2000个batch打印一次
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
scheduler.step()
print(f'Epoch {epoch+1} finished, lr: {scheduler.get_last_lr()[0]:.6f}')
训练技巧：使用学习率调度器（如CosineAnnealing）比固定学习率效果更好。初始学习率可以设大些（如0.1），让模型快速收敛。
5. 模型评估与可视化
5.1 测试集准确率计算
python复制testset = torchvision.datasets.CIFAR10(
root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(
testset, batch_size=100, shuffle=False, num_workers=2)
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data[0].to(device), data[1].to(device)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy on 10000 test images: {100 * correct / total:.2f}%')
5.2 混淆矩阵可视化
python复制from sklearn.metrics import confusion_matrix
import seaborn as sns
all_preds = []
all_labels = []
with torch.no_grad():
for data in testloader:
images, labels = data[0].to(device), data[1].to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)
all_preds.extend(predicted.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
cm = confusion_matrix(all_labels, all_preds)
plt.figure(figsize=(10,8))
sns.heatmap(cm, annot=True, fmt='d', xticklabels=classes, yticklabels=classes)
plt.xlabel('Predicted')
plt.ylabel('True')
plt.show()
6. 常见问题与解决方案
6.1 训练不收敛的可能原因
-
学习率设置不当：
- 现象：loss值波动大或几乎不变
- 解决：尝试不同学习率（如0.01, 0.001），或使用学习率finder工具
-
数据预处理问题：
- 现象：模型在训练集上表现也很差
- 解决：检查数据增强是否过度、归一化参数是否正确
-
模型结构缺陷：
- 现象：无论怎么调参效果都很差
- 解决：简化模型结构或换用经典结构（如ResNet18）
6.2 GPU内存不足的优化方法
- 减小batch size（如从64降到32）
- 使用梯度累积：
python复制accum_steps = 4 # 累积4个batch的梯度
for i, data in enumerate(trainloader):
inputs, labels = data[0].to(device), data[1].to(device)
outputs = net(inputs)
loss = criterion(outputs, labels)
loss = loss / accum_steps # 平均梯度
loss.backward()
if (i+1) % accum_steps == 0: # 每accum_steps步更新一次
optimizer.step()
optimizer.zero_grad()
- 使用混合精度训练：
python复制scaler = torch.cuda.amp.GradScaler()
for data in trainloader:
inputs, labels = data[0].to(device), data[1].to(device)
optimizer.zero_grad()
with torch.cuda.amp.autocast():
outputs = net(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
7. 进阶优化方向
当基础模型跑通后，可以考虑以下优化：
-
更强大的数据增强：
- 使用AutoAugment或RandAugment策略
- 尝试MixUp或CutMix等混合样本增强
-
模型结构改进：
- 换用EfficientNet、Vision Transformer等新架构
- 添加注意力机制模块
-
训练策略优化：
- 使用标签平滑（Label Smoothing）
- 实现知识蒸馏（用大模型指导小模型）
-
超参数调优：
- 使用Optuna或Ray Tune进行自动化超参数搜索
- 尝试不同的优化器（如AdamW）
我在实际项目中发现，对于CIFAR-10这样的小尺寸图片，适当减小模型规模（如减少ResNet的通道数）反而可能获得更好的效果，因为大模型容易过拟合。另外，在训练后期冻结前面的层、只微调最后几层，往往能带来1-2个百分点的精度提升。