实战：训练自己的图片分类器

这是什么：微调（fine-tuning） = 拿现成的预训练模型，换掉它"宣布答案"那一层，用自己的少量数据再训几轮。

为什么用：上一章借巨人的肩膀认出了 1000 种物体，这一次做自己的项目：微调一个只认识"猫、狗、兔子"的分类器。数据不用下载，手机相册就够了。

如何用：本章从"按类别分文件夹"开始，用 ImageFolder 自动标标签，再八二分训练集、验证集开训。

第一步：把图片按类别放进文件夹

准备每类几十到几百张图，按类别分文件夹——文件夹名就是标签：

data/
├── cat/      # 几十张猫图
├── dog/      # 几十张狗图
└── rabbit/   # 几十张兔子图

torchvision 的 ImageFolder 会自动扫描这个结构：cat 文件夹里的图自动标为类别 0……一行标注代码都不用写。

第二步：迁移学习微调

完整流程如下。小数据集在普通电脑上几分钟就能跑完（输出为示意）：

import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torchvision.models import resnet18, ResNet18_Weights
from torch.utils.data import DataLoader

# 1. 数据：ImageFolder 自动把"文件夹名"当标签
tf = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),          # 含归一化到 0~1
])
dataset = datasets.ImageFolder("data", transform=tf)
print(dataset.class_to_idx)         # 文件夹名 → 编号

# 2. 八二分成训练集 / 验证集，装进 DataLoader 分批喂
n_train = int(len(dataset) * 0.8)
train_set, val_set = torch.utils.data.random_split(
    dataset, [n_train, len(dataset) - n_train])
train_loader = DataLoader(train_set, batch_size=8, shuffle=True)
val_loader = DataLoader(val_set, batch_size=8)

# 3. 迁移学习：预训练 ResNet + 换一个新"头"
weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
model.fc = nn.Linear(model.fc.in_features, len(dataset.classes))

# 4. 只微调新换的头（预训练底层不动）
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

for epoch in range(5):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()

    # 每轮结束在验证集上看正确率
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            correct += (model(x).argmax(1) == y).sum().item()
            total += y.size(0)
    print(f"第 {epoch + 1} 轮：验证正确率 {correct / total:.0%}")

你会看到（输出为示意）：

{'cat': 0, 'dog': 1, 'rabbit': 2}
第 1 轮：验证正确率 68%
第 2 轮：验证正确率 81%
第 3 轮：验证正确率 88%
第 4 轮：验证正确率 92%
第 5 轮：验证正确率 94%

流程解读

零基础的算力建议

• CPU 就能跑：几十张图 + 只微调最后一层，普通笔记本 CPU 几分钟出结果。别被"深度学习必须显卡"吓住——那是指从零训大模型的场景
• 免费云端：数据量大、想跑快一点时，Kaggle、百度 AI Studio 等平台提供免费 GPU Notebook，注册就能用
• 从小做起：先把 3 个类别、每类 50 张跑通，再谈扩大规模

记忆口诀：图片按文件夹分，ImageFolder 自动标；底座来自预训练，换头微调就出师。

动手改

1. 把类别换成你感兴趣的图片（花卉、零食、手写数字都行），重新跑通全流程
2. 把 range(5) 改成 range(10)，观察验证正确率是继续涨还是开始波动
3. 写一个 predict(path) 函数：读入单张图片，打印预测的类别名（提示：复用上一章的预处理代码）

到这里，AI 板块全部更新完毕。感谢一路跟到毕业 🐸

图像识别毕业 🎉 AI 板块完结！回到首页看看其他板块，或去 Python 板块 巩固基础