用 fastai 训练一个几何图形分类器

fastai 提倡先训出能用的模型。这篇不下载 ImageNette，而是用 PIL 合成圆形/方形小数据集，再用 vision_learner(resnet18) + fine_tune 训练，并 export 出 pkl。 合成数据让你在无 GPU、无外部数据集授权问题的机器上完整走通官方训练 API。

先看清成品

• 真实训练：80 张训练 / 20 张验证，fine_tune 2 轮后 valid accuracy=1.0，导出 exports/shapes.pkl。
• make_shapes 会清空旧类别目录再写入，避免上次 n 更大时残留文件。
• 2 项 pytest 通过（含 1 轮 smoke 训练）。

合规提醒：只用你有权使用的图像数据；不要爬取或使用未授权的人脸/版权图片。load_learner 使用 pickle，只加载自己导出或可信来源的模型文件。本教程合成几何图形，不涉及真人数据。

运行说明：在 box 上用 CPU 版 PyTorch 2.14.1 + fastai 2.8.12 真实训练与推理（2026-10-23，UTC+8）。首次训练从 download.pytorch.org 下载官方 resnet18 权重约 45MB。未使用 GPU；准确率在合成圆形/方形任务上达到 1.0，不代表复杂真实数据集表现。

做完之后你会得到什么？

需要准备哪些环境？

关键事实

文件结构

95-fast-ai-A-src/
├── make_shapes.py
├── train_shapes.py
└── test_95a.py

第 1 步：合成数据

#!/usr/bin/env python3
"""Create a tiny synthetic shapes dataset (circle vs square) as PNG files."""
from __future__ import annotations

import argparse
import random
from pathlib import Path

from PIL import Image, ImageDraw


def draw_circle(size: int, rng: random.Random) -> Image.Image:
    img = Image.new("RGB", (size, size), (rng.randint(200, 255),) * 3)
    d = ImageDraw.Draw(img)
    m = rng.randint(4, size // 5)
    color = (rng.randint(0, 180), rng.randint(0, 180), rng.randint(0, 180))
    d.ellipse([m, m, size - m, size - m], fill=color)
    return img


def draw_square(size: int, rng: random.Random) -> Image.Image:
    img = Image.new("RGB", (size, size), (rng.randint(200, 255),) * 3)
    d = ImageDraw.Draw(img)
    m = rng.randint(4, size // 5)
    color = (rng.randint(0, 180), rng.randint(0, 180), rng.randint(0, 180))
    d.rectangle([m, m, size - m, size - m], fill=color)
    return img


def build(root: Path, n_per_class: int = 40, size: int = 64, seed: int = 0) -> Path:
    rng = random.Random(seed)
    for split, count in (("train", n_per_class), ("valid", max(2, n_per_class // 4))):
        for label, drawer in (("circle", draw_circle), ("square", draw_square)):
            d = root / split / label
            if d.exists():
                for oldf in d.glob("*"):
                    oldf.unlink()
            d.mkdir(parents=True, exist_ok=True)
            for i in range(count):
                drawer(size, rng).save(d / f"{label}_{i:03d}.png")
    return root


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--out", type=Path, default=Path("shapes"))
    p.add_argument("--n", type=int, default=40)
    args = p.parse_args(argv)
    build(args.out, n_per_class=args.n)
    print(f"wrote {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

第 2 步：训练

#!/usr/bin/env python3
"""Train a tiny image classifier with fastai on synthetic shapes.

Uses vision_learner(resnet18) + fine_tune. The first run downloads official
ImageNet weights from download.pytorch.org (~45MB); no third-party mirrors.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from make_shapes import build


def train(data: Path, epochs: int, export_path: Path) -> dict:
    from fastai.vision.all import ImageDataLoaders, Resize, accuracy, resnet18, vision_learner

    dls = ImageDataLoaders.from_folder(
        data,
        train="train",
        valid="valid",
        item_tfms=Resize(64),
        bs=16,
    )
    learn = vision_learner(dls, resnet18, metrics=accuracy, pretrained=True)
    # 1 frozen epoch + `epochs` unfrozen — classic fastai fine_tune
    learn.fine_tune(epochs, base_lr=1e-3)
    met = learn.recorder.values[-1] if learn.recorder.values else []
    acc = float(met[2]) if len(met) >= 3 else None
    # fastai resolves relative export paths under learn.path (= data root)
    export_path = export_path.resolve()
    export_path.parent.mkdir(parents=True, exist_ok=True)
    learn.export(export_path)
    return {
        "accuracy": acc,
        "export": str(export_path),
        "n_train": len(dls.train_ds),
        "n_valid": len(dls.valid_ds),
    }


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--data", type=Path, default=Path("shapes"))
    p.add_argument("--epochs", type=int, default=2)
    p.add_argument("--export", type=Path, default=Path("exports/shapes.pkl"))
    p.add_argument("--rebuild", action="store_true")
    p.add_argument("--n", type=int, default=40, help="images per class in train/")
    args = p.parse_args(argv)
    if args.rebuild or not (args.data / "train").exists():
        build(args.data, n_per_class=args.n)
    info = train(args.data, args.epochs, args.export)
    print(json.dumps(info, ensure_ascii=False, indent=2))
    if info["accuracy"] is not None and info["accuracy"] < 0.9:
        raise SystemExit(f"accuracy too low: {info['accuracy']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

第 3 步：测试

from pathlib import Path
from make_shapes import build

def test_build(tmp_path: Path):
    root = build(tmp_path / "shapes", n_per_class=8, size=32)
    assert (root / "train/circle").is_dir()
    assert len(list((root / "train/circle").glob("*.png"))) == 8
    assert len(list((root / "valid/square").glob("*.png"))) == 2

def test_train_smoke(tmp_path: Path):
    import importlib.util
    if importlib.util.find_spec("fastai") is None or importlib.util.find_spec("torch") is None:
        import pytest
        pytest.skip("fastai/torch not installed")
    from train_shapes import train
    data = build(tmp_path / "shapes", n_per_class=20, size=64)
    info = train(data, epochs=1, export_path=tmp_path / "m.pkl")
    assert info["n_train"] > 0
    assert Path(info["export"]).is_file()
    assert info["accuracy"] is None or info["accuracy"] >= 0.85

第 4 步：运行

$ python make_shapes.py --out shapes --n 40
wrote shapes
$ python train_shapes.py --data shapes --epochs 2 --export exports/shapes.pkl
# 首次会从 download.pytorch.org 下载 resnet18 权重约 45MB（官方源）
# 2026-10-23 12:38 UTC+8（box CPU）
epoch     train_loss  valid_loss  accuracy  time
0         1.051598    1.252751    0.550000  00:00   # fine_tune 冻结阶段
0         0.613625    0.308894    0.850000  00:00   # 解冻 epoch 1
1         0.396066    0.102641    1.000000  00:01   # 解冻 epoch 2
{
  "accuracy": 1.0,
  "export": "exports/shapes.pkl",
  "n_train": 80,
  "n_valid": 20
}
$ python -m pytest -q
2 passed

怎么检查结果？

• 结束后 accuracy ≥ 0.9，且 exports/shapes.pkl 存在。
• 换自己的数据时保持 train/valid/类别 目录约定。
• 第一次运行需能访问 download.pytorch.org。

调试示例

相对导出路径掉进数据集目录。 第一版 learn.export(Path("exports/shapes.pkl")) 实际写到 shapes/exports/...，因为 fastai 把相对路径解析到 learn.path。改为 export_path.resolve() 后文件落在项目 exports/。

从零训练 resnet18 不收敛。 合成任务虽简单，但 scratch 训练在短 epoch 内准确率停在 ~0.5。改用官方推荐的 pretrained + fine_tune 后两轮即到 1.0。教程默认走这条路径，并写明权重来源是 pytorch.org。

常见问题

模板规则

• 只用授权数据。
• export 用绝对路径。
• 预训练权重来自官方源。

项目资源

先在合成数据上把 API 跑顺，再换真实图片。