文章目录

一、前言

学深度学习最容易"看得懂原理、写不出代码"。本文带大家完成一个完整的 20 类食物图像分类项目，正好覆盖深度学习入门最关键的三个环节：

  1. 搭 CNN + 训练：自定义三层卷积网络，配合数据增强训练食物分类器；
  2. 调整学习率：用 torch.optim.lr_scheduler 在训练过程中自动调整学习率；
  3. 迁移学习：加载预训练 ResNet18，冻结主干、只替换最后一层，小数据也能快速出好结果；
  4. 加载最优模型测试：把训练好的 .pth 权重加载回来，在测试集上跑推理、算准确率。

二、环境准备

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 的训练固定五步，顺序不能乱：

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 迁移学习与冻结参数

小数据集从零训练容易过拟合，迁移学习的做法：

  1. 加载在 ImageNet 上预训练好的 ResNet18（weights=ResNet18_Weights.DEFAULT），它已经学会了通用的边缘、纹理特征；
  2. param.requires_grad = False 冻结主干，不再更新这些参数；
  3. 把最后的全连接层 fc 替换成 nn.Linear(512, 20)（ResNet18 原 fc 输入维度是 512），只训练新换的这一层；
  4. 优化器只接收 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 类。

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)

训练与优化：

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 迁移学习

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 加载最优模型测试

这一步不训练，只做推理：

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 训练 + 学习率调度

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 迁移学习

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 加载最优模型测试

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 权重离线推理 则是模型落地部署的最后一步。
# pytorch # cnn # 深度学习 # 迁移学习
Logo
亚马逊云科技技术品牌专区

更多推荐

cover

基于STM32的全屋环境智能监测与预警终端设计

avatar 亚马逊云科技技术品牌专区

Keras：面向人类的深度学习库

Keras3支持多后端（TensorFlow、PyTorch、JAX），通过设置环境变量切换，兼容指定版本。pip安装后可用Sequential API快速构建CNN模型，如MNIST分类示例，实

avatar 亚马逊云科技技术品牌专区

AI行业日报｜2026-09-25：5个热点事件

本期行业动态聚焦于端侧算力重构、自主智能体工作流、开发工具经济模型、网络隔离架构设计与公共部门合规审查。以下结合一线技术场景进行结构化梳理，供 AI 研究员与软件工程团队参考。 PrismML 正将其轻量级语言模型适配至搭载高通处理器的智能眼镜终端 1 。该平台的核心战略指向开源权重模型的本地化运行，旨在最大化挖掘穿戴设备现有算力的利用效率 1 。对于嵌入式与边缘计算开发者而言，此类微型模型向低功

avatar 亚马逊云科技技术品牌专区