用 fastai 导出分类模型并在推理端加载预测

训练机和推理机常常分离。fastai 的 export / load_learner 把预处理与模型打包；官方也提醒 pickle 不安全，只加载可信文件。 这篇做 infer.py：加载 A 篇导出的 pkl，对单张图输出标签与概率。

先看清成品

• 对 valid/circle 样本推理得到 label=circle，概率约 0.98。
• pytest 完整 roundtrip（短训+推理）通过。
• 运行时出现官方 pickle 安全警告，已在文中说明。

合规提醒：只用你有权使用的图像数据；不要爬取或使用未授权的人脸/版权图片。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-C-src/
├── infer.py
└── test_95c.py

第 1 步：推理脚本

#!/usr/bin/env python3
"""Load a fastai export.pkl and predict a single image.

Official note: Learner.export pickles function *names*, so any custom label_func
must be importable in the inference environment. This script keeps labels as
folder names (from_folder) so no custom function is needed.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path


def predict(export_path: Path, image: Path) -> dict:
    from fastai.vision.all import load_learner, PILImage

    learn = load_learner(export_path)
    pred, idx, probs = learn.predict(PILImage.create(image))
    return {
        "label": str(pred),
        "index": int(idx),
        "probs": {str(c): float(probs[i]) for i, c in enumerate(learn.dls.vocab)},
    }


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("export", type=Path)
    p.add_argument("image", type=Path)
    args = p.parse_args(argv)
    print(json.dumps(predict(args.export, args.image), ensure_ascii=False, indent=2))
    return 0


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

第 2 步：测试

from pathlib import Path
import importlib.util
import pytest

def test_infer_roundtrip(tmp_path: Path):
    if importlib.util.find_spec("fastai") is None or importlib.util.find_spec("torch") is None:
        pytest.skip("fastai/torch not installed")
    # reuse 95A builders via path hack
    import sys
    sys.path.insert(0, str(Path("/tmp/p/95A")))
    from make_shapes import build
    from train_shapes import train
    from infer import predict

    data = build(tmp_path / "shapes", n_per_class=16, size=32)
    exp = tmp_path / "m.pkl"
    train(data, epochs=2, export_path=exp)
    img = next((data / "valid" / "circle").glob("*.png"))
    out = predict(exp, img)
    assert "label" in out
    assert "circle" in out["probs"] and "square" in out["probs"]

第 3 步：运行

$ python infer.py exports/shapes.pkl shapes/valid/circle/circle_000.png
{
  "label": "circle",
  "index": 0,
  "probs": {
    "circle": 0.9832153916358948,
    "square": 0.016784552484750748
  }
}
$ python -m pytest -q
1 passed

# 实际概率以你的导出为准；实录（2026-10-23 UTC+8）：
█

 |----------------------------------------| 0.00% [0/1 00:00<?]
 |████████████████████████████████████████| 100.00% [1/1 00:00<00:00]
                                                                     

                                                                     
{
  "label": "circle",
  "index": 0,
  "probs": {
    "circle": 0.9832153916358948,
    "square": 0.016784552484750748
  }
}

怎么检查结果？

• 输出含 label 与两类概率且和约为 1。
• 只加载自己 export 的 pkl。
• 推理环境需安装同主版本的 fastai/torch。

调试示例

pickle 警告不是噪音。 fastai 2.8 会明确警告 load_learner 使用 insecure pickle。生产应：① 只接受内部构建产物；② 考虑只保存 weights 并用 Learner.load；③ 不要加载聊天里陌生人发来的 pkl。

常见问题

模板规则

• 只加载可信 pkl。
• 推理依赖与训练对齐。
• 记录 vocab。

项目资源

导出是交付物；pickle 是信任边界。