【深度学习入门】PyTorch 食物图像分类三连：CNN 训练、学习率调度与 ResNet18 迁移学习

文章目录

一、前言
二、环境准备
三、必备知识补充
- [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 权重离线推理 则是模型落地部署的最后一步。