第 P5 周：PyTorch 实现运动鞋识别（动态学习率）

一、本周学习内容概览

本周 P5 的核心任务是：使用 PyTorch 搭建 CNN 对 adidas / nike 运动鞋图片进行二分类识别，并重点学习和实践动态学习率调度策略。

我这次主要做了以下几件事：

  • 数据组织：数据集 46-data.zip 解压后为 5-data/，内部已按 train/ 和 test/ 划分好，每个子目录下有 adidas/ 和 nike/ 两类图片。
  • 数据加载与预处理：使用 datasets.ImageFolder 分别读取训练集和测试集，通过 transforms 将图片统一缩放到 224×224 并按 ImageNet 常用的 mean/std 做标准化。
  • 模型构建：搭建了带有 BatchNorm 和 Dropout 的 CNN（Model），包含四层卷积 + 两次池化 + 全连接层，输出 2 维 logits。
  • 动态学习率（重点）：使用 torch.optim.lr_scheduler.StepLR，每 5 个 epoch 将学习率衰减为原来的 0.5 倍，替代了手动写 adjust_learning_rate 的方式。
  • 训练与评估：训练 40 个 epoch，记录每轮的 Train/Test Acc/Loss 和当前学习率，并在训练过程中保存最佳模型权重。
  • 最终结果：训练集准确率约 99.6%，测试集最佳准确率约 82.9%。

二、必要代码与关键点说明

1. 环境与设备检查

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torchvision import datasets
import os, PIL, pathlib

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

2. 数据集目录结构与路径处理

数据集解压后结构如下（已按 train/test 划分好）：

5-data/
  train/
    adidas/
      img001.jpg ...
    nike/
      img101.jpg ...
  test/
    adidas/
      img201.jpg ...
    nike/
      img301.jpg ...

使用 pathlib.Path 处理路径，并从 ImageFolder 直接获取类别名：

from pathlib import Path

data_dir = Path("5-data")
data_paths = list(data_dir.glob("*"))
print([p.name for p in data_paths])  # ['test', 'train']

3. 数据加载与预处理

train_transforms = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

test_transform = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder("5-data/train", transform=train_transforms)
test_dataset  = datasets.ImageFolder("5-data/test",  transform=test_transform)

classNames = train_dataset.classes   # ['adidas', 'nike']
train_dataset.class_to_idx           # {'adidas': 0, 'nike': 1}

batch_size = 32
train_dl = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True,  num_workers=1)
test_dl  = torch.utils.data.DataLoader(test_dataset,  batch_size=batch_size, shuffle=False, num_workers=1)

4. CNN 模型结构

import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, num_classes):
        super(Model, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(3, 12, 5), nn.BatchNorm2d(12), nn.ReLU())   # -> 12×220×220
        self.conv2 = nn.Sequential(nn.Conv2d(12, 12, 5), nn.BatchNorm2d(12), nn.ReLU())  # -> 12×216×216
        self.pool3 = nn.Sequential(nn.MaxPool2d(2))                                       # -> 12×108×108
        self.conv4 = nn.Sequential(nn.Conv2d(12, 24, 5), nn.BatchNorm2d(24), nn.ReLU())  # -> 24×104×104
        self.conv5 = nn.Sequential(nn.Conv2d(24, 24, 5), nn.BatchNorm2d(24), nn.ReLU())  # -> 24×100×100
        self.pool6 = nn.Sequential(nn.MaxPool2d(2))                                       # -> 24×50×50
        self.dropout = nn.Sequential(nn.Dropout(0.2))
        self.fc = nn.Sequential(nn.Linear(24 * 50 * 50, num_classes))

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.pool3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.pool6(x)
        x = self.dropout(x)
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x

model = Model(num_classes=len(train_dataset.classes)).to(device)

5. 损失函数、优化器与动态学习率（重点）

loss_fn = nn.CrossEntropyLoss()
learn_rate = 1e-3
optimizer = torch.optim.SGD(model.parameters(), lr=learn_rate, momentum=0.9, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
  • StepLR：每 5 个 epoch 学习率乘以 gamma=0.5（即减半）。
  • 学习率变化轨迹：1e-3 → 5e-4 → 2.5e-4 → 1.25e-4 → ...

三、训练与测试流程

1. 训练函数

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    train_loss, train_acc = 0.0, 0.0

    for X, y in dataloader:
        X, y = X.to(device), y.to(device)
        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_acc += (pred.argmax(1) == y).type(torch.float).sum().item()
        train_loss += loss.item()

    return train_acc / size, train_loss / num_batches

2. 测试函数

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, test_acc = 0.0, 0.0

    with torch.no_grad():
        for imgs, target in dataloader:
            imgs, target = imgs.to(device), target.to(device)
            target_pred = model(imgs)
            loss = loss_fn(target_pred, target)
            test_loss += loss.item()
            test_acc += (target_pred.argmax(1) == target).type(torch.float).sum().item()

    return test_acc / size, test_loss / num_batches

3. 主训练循环（含最佳模型保存）

epochs = 40
train_loss, train_acc = [], []
test_loss, test_acc = [], []
best_acc = 0.0
best_path = "best_model_p5.pth"

for epoch in range(epochs):
    model.train()
    epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, optimizer)

    model.eval()
    epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn)

    scheduler.step()

    train_acc.append(epoch_train_acc)
    train_loss.append(epoch_train_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)

    if epoch_test_acc > best_acc:
        best_acc = epoch_test_acc
        torch.save(model.state_dict(), best_path)

    lr = optimizer.param_groups[0]['lr']
    template = "Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}"
    print(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss,
                          epoch_test_acc*100, epoch_test_loss, lr))

print("Done")
print(f"Best Test Acc: {best_acc*100:.2f}% | Saved to: {best_path}")

在本次实验中，40 个 epoch 后的代表性结果为：

请添加图片描述

  • 训练集：Train_acc ≈ 99.6%，Train_loss ≈ 0.04
  • 测试集最佳：Test_acc ≈ 82.9%（出现在 Epoch 10–12 附近）
  • 学习率变化：从 1e-3 逐步衰减到 3.91e-6

四、结果可视化与分析

1. 绘制训练 / 测试曲线

import matplotlib.pyplot as plt
from datetime import datetime

epochs_range = range(epochs)
plt.figure(figsize=(12, 3))

plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc, label="Train Acc")
plt.plot(epochs_range, test_acc, label="Test Acc")
plt.legend(loc="lower right")
plt.title("Training and Validation Accuracy")
plt.xlabel(datetime.now())

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss, label="Train Loss")
plt.plot(epochs_range, test_loss, label="Test Loss")
plt.legend(loc="upper right")
plt.title("Training and Validation Loss")
plt.tight_layout()
plt.show()

请添加图片描述

从曲线可以看到：训练集准确率在 10 个 epoch 后就达到 95%+ 并持续上升到接近 100%，而测试集在 82% 左右就趋于平稳，说明模型存在一定过拟合。

2. 指定图片预测

from PIL import Image
from pathlib import Path
import random

idx_to_class = {v: k for k, v in train_dataset.class_to_idx.items()}

test_root = Path("5-data/test")
img_candidates = list(test_root.glob("*/*.jpg")) + list(test_root.glob("*/*.png"))
img_path = random.choice(img_candidates)

img = Image.open(img_path).convert("RGB")
x = test_transform(img)
x = torch.unsqueeze(x, dim=0).to(device)

model.eval()
with torch.no_grad():
    pred_idx = model(x).argmax(dim=1).item()

print(f"图片: {img_path}")
print(f"预测类别: {idx_to_class[pred_idx]}")

请添加图片描述

五、保存并加载模型

# 保存最后一轮权重
torch.save(model.state_dict(), "last_model_p5.pth")

# 加载最佳权重
best_model = Model(num_classes=len(train_dataset.classes)).to(device)
best_model.load_state_dict(torch.load("best_model_p5.pth", map_location=device))
best_model.eval()

六、动态学习率知识总结

本次重点学习了三种官方学习率调度器：

1) StepLR（本次实际使用）

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

每 step_size 个 epoch，学习率乘 gamma。

2) LambdaLR（自定义衰减函数）

lambda1 = lambda epoch: 0.92 ** (epoch // 2)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)

3) MultiStepLR（指定里程碑衰减）

scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer, milestones=[10, 20, 30], gamma=0.1
)

在指定的 epoch 节点上衰减，适合已知训练曲线拐点的场景。

七、个人总结

  • 环境方面：本次实验在 WSL2 + Linux 环境下完成，Python 版本为 3.11.13，PyTorch 为 2.5.1+cu121，Torchvision 为 0.20.1+cu121，GPU 为 NVIDIA GeForce RTX 4060 Ti。一开始在路径与类别数上再次踩了坑（split("\\")、全局变量 classeNames 拼写不一致导致 out_features=0 和 CUDA assert），重启内核并统一使用 train_dataset.classes 后解决。
  • 数据与模型方面：本次使用的是 adidas vs nike 运动鞋图片数据集，数据量不大。自定义的 Model（四层卷积 + BN + Dropout + 全连接）在 40 个 epoch 后训练集准确率达到 99.6%，但测试集最高只到约 82.9%，过拟合较为明显。
  • 动态学习率（本次重点）：体会到了动态调整学习率的必要性——初期用较大的 lr=1e-3 快速收敛，后期通过 StepLR 逐步减小学习率来稳定训练。相比固定 lr=1e-4 的 SGD，收敛速度和最终精度都有所改善。还了解了 LambdaLR、MultiStepLR 等更灵活的调度方式。
  • 最佳模型保存（拔高）：在训练循环中加入了对测试集准确率的跟踪，只有当新的 test_acc 超过历史最佳时才保存权重到 best_model_p5.pth，避免最终使用过拟合的最后一轮模型。
  • 后续改进方向：测试集准确率要达到 84%–86%，可以尝试加入 RandomHorizontalFlip 等数据增强、使用 Adam 优化器、或引入预训练 ResNet18 做迁移学习，同时适度增加 Dropout 强度来控制过拟合。
# 学习
Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

AtomGit 首发开源｜光轮智能发布全球首个十万小时级全模态人类行为开源数据集 EgoSuite-Open100K

avatar AtomGit开源社区
cover

构筑赋能金融行业应用 Agent｜20 强决战北京，一场金融 Agent 的硬核答卷

avatar AtomGit开源社区
cover

AtomCode 全新上线万亿代码模型 LongCat-2.0。限时 7 天免费体验！

avatar AtomGit开源社区