文章目录
- 一、前言
- 二、环境准备
- 三、必备知识补充
-
- [3.1 卷积层输出尺寸怎么算](#3.1 卷积层输出尺寸怎么算)
- [3.2 自定义 Dataset 三件套](#3.2 自定义 Dataset 三件套)
- [3.3 数据增强与 ImageNet 归一化](#3.3 数据增强与 ImageNet 归一化)
- [3.4 训练循环五步法](#3.4 训练循环五步法)
- [3.5 学习率调度器(lr_scheduler)](#3.5 学习率调度器(lr_scheduler))
- [3.6 迁移学习与冻结参数](#3.6 迁移学习与冻结参数)
- [3.7 保存与加载权重](#3.7 保存与加载权重)
- 四、代码精讲
-
- [4.1 数据增强 + CNN 训练 + 学习率调度](#4.1 数据增强 + CNN 训练 + 学习率调度)
- [4.2 ResNet18 迁移学习](#4.2 ResNet18 迁移学习)
- [4.3 加载最优模型测试](#4.3 加载最优模型测试)
-
- [4.4.1 CNN 训练 + 学习率调度](#4.4.1 CNN 训练 + 学习率调度)
- [4.4.2 ResNet18 迁移学习](#4.4.2 ResNet18 迁移学习)
- [4.4.3 加载最优模型测试](#4.4.3 加载最优模型测试)
- 五、关键技巧与常见问题
- 六、总结
一、前言
学深度学习最容易"看得懂原理、写不出代码"。本文带大家完成一个完整的 20 类食物图像分类项目,正好覆盖深度学习入门最关键的三个环节:
- 搭 CNN + 训练:自定义三层卷积网络,配合数据增强训练食物分类器;
- 调整学习率 :用
torch.optim.lr_scheduler在训练过程中自动调整学习率; - 迁移学习:加载预训练 ResNet18,冻结主干、只替换最后一层,小数据也能快速出好结果;
- 加载最优模型测试 :把训练好的
.pth权重加载回来,在测试集上跑推理、算准确率。
二、环境准备
bash
pip install torch torchvision
pip install pillow numpy
说明:
- 代码会自动选择设备:有 NVIDIA 显卡用
cuda,Apple 芯片 Mac 用mps,都没有就退回cpu; - 数据集是自定义格式:
train.txt/test.txt每行一条,内容为图片路径 标签(空格分隔,标签为 0--19 的数字字符串); - 训练好的权重文件命名为
best_model.pth,本文通过torch.load加载。
三、必备知识补充
3.1 卷积层输出尺寸怎么算
卷积层参数 Conv2d(in_channels, out_channels, kernel_size, stride, padding),输出特征图尺寸公式:
输出尺寸 = floor((输入尺寸 - 卷积核尺寸 + 2×填充) / 步长) + 1
代入本项目:输入 256×256,卷积核 5、步长 1、填充 2:(256 - 5 + 4)/1 + 1 = 256,尺寸不变 ;每经过一次 MaxPool2d(2),宽高都减半。所以数据形状变化为:
| 层 | 输出尺寸(高×宽×通道) |
|---|---|
| 输入 | 256 × 256 × 3 |
| conv1(卷积+ReLU+池化) | 128 × 128 × 16 |
| conv2(两次卷积+ReLU+池化) | 64 × 64 × 32 |
| conv3(卷积+ReLU,不池化) | 64 × 64 × 128 |
| 展平后接全连接层 | 128 × 64 × 64 = 524288 → 20 类 |
这也是 nn.Linear(128 * 64 * 64, 20) 数字的由来------改输入尺寸或网络结构后,这里必须跟着改,否则会报维度不匹配错误。
3.2 自定义 Dataset 三件套
PyTorch 读取自定义数据集只要继承 Dataset 并实现三个方法:
__init__:记录图片路径列表、标签列表和预处理方法;__len__:返回样本总数;__getitem__(idx):按索引读一张图、做预处理、返回(图像张量, 标签张量)。
DataLoader 再把这些样本按 batch_size=64 打包、shuffle=True 打乱,训练时按批取数。
3.3 数据增强与 ImageNet 归一化
训练集用了这些随机增强(验证集一律不用随机操作,保证评估稳定):
| 操作 | 作用 |
|---|---|
Resize([256,256]) |
统一缩放到 256×256 |
RandomRotation(45) |
随机旋转 ±45° |
CenterCrop(256) |
中心裁剪 |
RandomHorizontalFlip(0.5) |
50% 概率水平翻转 |
RandomVerticalFlip(0.5) |
50% 概率垂直翻转 |
ColorJitter(...) |
随机调亮度/对比度/饱和度/色调 |
RandomGrayscale(0.1) |
10% 概率转灰度 |
最后两步容易被忽略:
ToTensor():把 PIL 图像(0--255)转成张量并缩放到[0,1],维度从[H,W,C]变成[C,H,W];Normalize(mean, std):用 ImageNet 统计的均值[0.485,0.456,0.406]、方差[0.229,0.224,0.225]做标准化。用了预训练模型(如 ResNet18)时必须用这套数值,否则模型性能会明显下降。
3.4 训练循环五步法
每个 batch 的训练固定五步,顺序不能乱:
python
pred = model.forward(X) # 1. 前向传播
loss = loss_fn(pred, y) # 2. 计算损失
optimizer.zero_grad() # 3. 梯度清零(PyTorch梯度会累加,必须清)
loss.backward() # 4. 反向传播,自动求梯度
optimizer.step() # 5. 按梯度更新权重
- 损失函数用
CrossEntropyLoss():它内部已经包含 softmax,直接喂模型输出的原始 logits 即可,不要再自己加 softmax; model.train()/model.eval()切换训练/评估模式;- 推理时用
with torch.no_grad():关闭梯度计算,省显存、更快; pred.argmax(1)在类别维度取最大分数对应的索引,即预测类别。
3.5 学习率调度器(lr_scheduler)
学习率太大容易震荡不收敛,太小又收敛慢,所以训练中通常要"先大后小"。本文件里出现的几种调度器:
| 调度器 | 行为 |
|---|---|
ReduceLROnPlateau |
监视某个指标(这里是训练 loss),连续 10 轮不下降就把学习率 ×0.1;适合不知道何时该衰减时 |
StepLR(step_size=5, gamma=0.5) |
每 5 个 epoch 学习率 ×0.5 |
MultiStepLR(milestones=[3,5,7], gamma=0.1) |
在第 3、5、7 轮各衰减一次 |
ExponentialLR(gamma=0.1) |
每轮学习率都 ×0.1(衰减很快) |
CosineAnnealingLR(T_max=5) |
按余弦曲线把学习率平滑降到 eta_min |
调用注意:ReduceLROnPlateau.step(指标值) 要传入监视的数值,且通常每个 epoch 结束后调用一次 ;StepLR 等则每个 epoch 调用一次无参 scheduler.step()。
3.6 迁移学习与冻结参数
小数据集从零训练容易过拟合,迁移学习的做法:
- 加载在 ImageNet 上预训练好的 ResNet18(
weights=ResNet18_Weights.DEFAULT),它已经学会了通用的边缘、纹理特征; param.requires_grad = False冻结主干,不再更新这些参数;- 把最后的全连接层
fc替换成nn.Linear(512, 20)(ResNet18 原 fc 输入维度是 512),只训练新换的这一层; - 优化器只接收
requires_grad=True的参数列表,训练又快又稳。
3.7 保存与加载权重
- 保存:
torch.save(model.state_dict(), 'best.pth'),存的是参数字典,不是整个模型; - 加载:先
model = CNN()搭好相同结构,再model.load_state_dict(torch.load('best_model.pth')),最后model.eval()才能推理。
四、代码精讲
4.1 数据增强 + CNN 训练 + 学习率调度
网络结构:三个卷积块,最终输出 20 类。
python
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=2),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2)
)
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2),
nn.ReLU(),
nn.Conv2d(32, 32, 5, 1, 2),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.conv3 = nn.Sequential(
nn.Conv2d(32, 128, 5, 1, 2),
nn.ReLU()
)
self.out = nn.Linear(128 * 64 * 64, 20)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1) # 展平,保留 batch 维度
return self.out(x)
训练与优化:
python
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', factor=0.1, patience=10, min_lr=0)
epochs = 10
for t in range(epochs):
print(f"Epoch {t+1}")
train(train_dataloader, model, loss_fn, optimizer)
test(test_dataloader, model, loss_fn)
训练函数就是前面讲的"五步法";测试函数里 model.eval() + torch.no_grad(),并用 pred.argmax(1) == y 统计正确样本数。
4.2 ResNet18 迁移学习
python
resnet_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for param in resnet_model.parameters():
param.requires_grad = False # 冻结主干
in_features = resnet_model.fc.in_features # 512
resnet_model.fc = nn.Linear(in_features, 20) # 只换最后一层
params_to_update = [p for p in resnet_model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(params_to_update, lr=0.001)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
每个 epoch:训练一轮 → scheduler.step() 调整学习率 → 在测试集上评估,并记录历史最优准确率 best_acc。
4.3 加载最优模型测试
这一步不训练,只做推理:
python
model = CNN().to(device)
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
with torch.no_grad():
for X, y in test_dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
result.append(pred.argmax(1).item())
labels.append(y.item())
# 统计准确率
correct = sum(p == l for p, l in zip(result, labels))
print('总数', len(labels), '正确数', correct, '正确率', correct / len(labels))
注意:这里的预处理只用
Resize + ToTensor,没有 RandomRotation 等随机增强,也没有 Normalize------推理阶段必须与训练时的预处理口径保持一致(如果模型训练时用了 Normalize,测试时也必须加上同一组均值方差)。
4.4.1 CNN 训练 + 学习率调度
python
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
from torchvision import transforms
# 数据增强与预处理
data_transforms = {
'trainda': transforms.Compose([
transforms.Resize([256, 256]),
transforms.RandomRotation(45),
transforms.CenterCrop(256),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
transforms.RandomGrayscale(p=0.1),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'valid': transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
class food_dataset(Dataset):
def __init__(self, file_path, transform=None):
self.file_path = file_path
self.imgs = []
self.labels = []
self.transform = transform
with open(self.file_path) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for img_path, label in samples:
self.imgs.append(img_path)
self.labels.append(label)
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
image = Image.open(self.imgs[idx])
if self.transform:
image = self.transform(image)
label = self.labels[idx]
label = torch.from_numpy(np.array(label, dtype=np.int64))
return image, label
training_data = food_dataset(file_path='./train.txt', transform=data_transforms['trainda'])
test_data = food_dataset(file_path='./test.txt', transform=data_transforms['valid'])
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=2),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2)
)
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2),
nn.ReLU(),
nn.Conv2d(32, 32, 5, 1, 2),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.conv3 = nn.Sequential(
nn.Conv2d(32, 128, 5, 1, 2),
nn.ReLU()
)
self.out = nn.Linear(128 * 64 * 64, 20)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1)
return self.out(x)
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
model = CNN().to(device)
def train(dataloader, model, loss_fn, optimizer):
model.train()
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = loss.item()
scheduler.step(loss) # 监视 loss,连续不下降则降低学习率
print(f'loss: {loss:>7f}')
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(f"Test result: \n Accuracy: {(100*correct)}%, Avg loss: {test_loss}")
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', factor=0.1, patience=10, threshold=0.0001, min_lr=0)
epochs = 10
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
print("Done!")
test(test_dataloader, model, loss_fn)
# 训练完成后保存权重
torch.save(model.state_dict(), 'best_model.pth')
4.4.2 ResNet18 迁移学习
python
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
from torchvision import transforms
import torchvision.models as models
# 加载预训练 ResNet18,冻结主干
resnet_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for param in resnet_model.parameters():
param.requires_grad = False
in_features = resnet_model.fc.in_features # 512
resnet_model.fc = nn.Linear(in_features, 20) # 替换为 20 类输出
params_to_update = [p for p in resnet_model.parameters() if p.requires_grad]
# 数据增强(与 4.4.1 一致)
data_transforms = {
'trainda': transforms.Compose([
transforms.Resize([256, 256]),
transforms.RandomRotation(45),
transforms.CenterCrop(256),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
transforms.RandomGrayscale(p=0.1),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'valid': transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
class food_dataset(Dataset):
def __init__(self, file_path, transform=None):
self.file_path = file_path
self.imgs = []
self.labels = []
self.transform = transform
with open(self.file_path) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for img_path, label in samples:
self.imgs.append(img_path)
self.labels.append(label)
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
image = Image.open(self.imgs[idx])
if self.transform:
image = self.transform(image)
label = self.labels[idx]
label = torch.from_numpy(np.array(label, dtype=np.int64))
return image, label
training_data = food_dataset(file_path='./train.txt', transform=data_transforms['trainda'])
test_data = food_dataset(file_path='./test.txt', transform=data_transforms['valid'])
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
print(f'Using {device} device')
model = resnet_model.to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params_to_update, lr=0.001)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
def train(dataloader, model, loss_fn, optimizer):
model.train()
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
best_acc = 0
acc_s, loss_s = [], []
def test(dataloader, model, loss_fn):
global best_acc
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(f"Test result: \n Accuracy: {(100*correct)}%, Avg loss: {test_loss}")
acc_s.append(correct)
loss_s.append(test_loss)
if correct > best_acc:
best_acc = correct
torch.save(model.state_dict(), 'best_model.pth') # 保存最优轮次权重
epochs = 10
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
scheduler.step()
test(test_dataloader, model, loss_fn)
print("最优训练结果为:", best_acc)
4.4.3 加载最优模型测试
python
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
from torchvision import transforms
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 16, 5, 1, 2),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2),
nn.ReLU(),
nn.Conv2d(32, 32, 5, 1, 2),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.conv3 = nn.Sequential(
nn.Conv2d(32, 128, 5, 1, 2),
nn.ReLU()
)
self.out = nn.Linear(128 * 64 * 64, 20)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1)
return self.out(x)
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
# 搭建与训练时相同的结构,再加载权重
model = CNN().to(device)
model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=True))
model.eval()
# 注意:预处理要与训练时一致,这里只用 Resize + ToTensor
data_transforms = transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
])
class food_dataset(Dataset):
def __init__(self, file_path, transform=None):
self.file_path = file_path
self.imgs = []
self.labels = []
self.transform = transform
with open(self.file_path) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for img_path, label in samples:
self.imgs.append(img_path)
self.labels.append(label)
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
image = Image.open(self.imgs[idx])
if self.transform:
image = self.transform(image)
label = self.labels[idx]
label = torch.from_numpy(np.array(label, dtype=np.int64))
return image, label
test_data = food_dataset(file_path='./test.txt', transform=data_transforms)
test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False)
result, labels = [], []
with torch.no_grad():
for X, y in test_dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
result.append(pred.argmax(1).item())
labels.append(y.item())
print('预测值:', result)
print('真实值:', labels)
correct = sum(p == l for p, l in zip(result, labels))
print('正确率', correct / len(labels))
五、关键技巧与常见问题
| 问题 | 说明 |
|---|---|
shape of input tensor is different / 维度不匹配 |
输入尺寸或网络结构变了,Linear 的输入维度要按 4.1 节重新算 |
| 准确率一直上不去 | 检查训练/测试预处理是否一致;是否漏了 Normalize;学习率是否合适 |
| 预训练模型效果差 | 几乎一定是没按 ImageNet 均值方差做 Normalize |
| GPU 显存不够 | 减小 batch_size;推理本来就该用 no_grad |
ReduceLROnPlateau 不生效 |
它要在每轮结束后调用 scheduler.step(指标值),不要在每个 batch 里频繁调用 |
| 加载权重报 key 不匹配 | 模型结构必须和保存时完全一致,且先实例化 CNN() 再 load_state_dict |
| 训练时 loss 不降反升 | 检查是否忘了 optimizer.zero_grad();学习率可能过大 |
小建议:
torch.load在新版 PyTorch 中建议写成torch.load('best_model.pth', map_location=device, weights_only=True),兼容性和安全性更好。
六、总结
本文串起了深度学习图像分类的最小完整闭环:
- 自定义 CNN 理解卷积、池化、全连接的数据流;
- 数据增强 + 学习率调度 是提升泛化、加速收敛的标准手段;
- ResNet18 迁移学习 用很少的数据和算力就能拿到不错的精度;
- 加载
.pth权重离线推理 则是模型落地部署的最后一步。