用 Python 检查 fastai 图像数据集目录结构

ImageDataLoaders.from_folder 对目录很敏感：train/valid 类别不一致、某类空文件夹，都会在训练时才爆。 这篇做 check_dataset.py，在训练前检查布局、扩展名和空文件。

先看清成品

• 合法 shapes 布局通过；valid 多出 triangle 类时报类别不一致。
• 3 项 pytest 通过。

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

第 1 步：检查脚本

#!/usr/bin/env python3
"""Check an ImageNet-style folder layout that ImageDataLoaders.from_folder expects."""
from __future__ import annotations

import argparse
import sys
from collections import Counter
from pathlib import Path

IMG_EXT = {".png", ".jpg", ".jpeg", ".bmp", ".webp", ".gif"}


def check(root: Path, train_name="train", valid_name="valid") -> list[str]:
    errs = []
    train = root / train_name
    valid = root / valid_name
    if not train.is_dir():
        errs.append(f"缺少 {train_name}/ 目录")
    if not valid.is_dir():
        # from_folder can use valid_pct instead, but this checker expects explicit valid/
        errs.append(f"缺少 {valid_name}/ 目录（若打算用 valid_pct 可忽略）")
        return errs
    train_labels = {p.name for p in train.iterdir() if p.is_dir()} if train.is_dir() else set()
    valid_labels = {p.name for p in valid.iterdir() if p.is_dir()}
    if not train_labels:
        errs.append("train 下没有类别子目录")
    if train_labels != valid_labels:
        errs.append(f"train 与 valid 类别不一致: train={sorted(train_labels)} valid={sorted(valid_labels)}")
    for split, base in ((train_name, train), (valid_name, valid)):
        if not base.is_dir():
            continue
        for label_dir in base.iterdir():
            if not label_dir.is_dir():
                continue
            imgs = [p for p in label_dir.iterdir() if p.suffix.lower() in IMG_EXT]
            if not imgs:
                errs.append(f"{split}/{label_dir.name}/ 没有图片")
            # empty files
            for p in imgs:
                if p.stat().st_size == 0:
                    errs.append(f"空文件: {p.relative_to(root)}")
    return errs


def summary(root: Path) -> str:
    lines = []
    for split in ("train", "valid"):
        base = root / split
        if not base.is_dir():
            continue
        counts = Counter()
        for label_dir in sorted(p for p in base.iterdir() if p.is_dir()):
            n = sum(1 for p in label_dir.iterdir() if p.suffix.lower() in IMG_EXT)
            counts[label_dir.name] = n
        lines.append(f"{split}: " + ", ".join(f"{k}={v}" for k, v in counts.items()))
    return "\n".join(lines)


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("root", type=Path)
    args = p.parse_args(argv)
    errs = check(args.root)
    print(summary(args.root) or "(空)")
    if errs:
        print(f"✗ {len(errs)} 个问题")
        for e in errs:
            print(f"  - {e}")
        return 1
    print("✓ 目录结构符合 ImageDataLoaders.from_folder(train/valid) 约定")
    return 0


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

第 2 步：测试

from pathlib import Path
from PIL import Image
from check_dataset import check, main

def _mk(root: Path, labels=("a", "b"), n=2):
    for split in ("train", "valid"):
        for lab in labels:
            d = root / split / lab
            d.mkdir(parents=True)
            for i in range(n):
                Image.new("RGB", (8, 8), (i * 40, 0, 0)).save(d / f"{i}.png")

def test_ok(tmp_path: Path):
    root = tmp_path / "d"
    _mk(root)
    assert check(root) == []
    assert main([str(root)]) == 0

def test_mismatch(tmp_path: Path):
    root = tmp_path / "d"
    _mk(root)
    (root / "valid" / "c").mkdir()
    Image.new("RGB", (8, 8)).save(root / "valid" / "c" / "0.png")
    errs = check(root)
    assert any("不一致" in e for e in errs)

def test_empty_class(tmp_path: Path):
    root = tmp_path / "d"
    _mk(root)
    for p in (root / "train" / "a").glob("*"):
        p.unlink()
    errs = check(root)
    assert any("没有图片" in e for e in errs)

第 3 步：运行

# 2026-10-23 12:38 UTC+8
$ python check_dataset.py ok_shapes/
train: circle=8, square=8
valid: circle=2, square=2
✓ 目录结构符合 ImageDataLoaders.from_folder(train/valid) 约定
$ python check_dataset.py bad_shapes/
train: circle=8, square=8
valid: circle=2, square=2, triangle=1
✗ 1 个问题
  - train 与 valid 类别不一致: train=['circle', 'square'] valid=['circle', 'square', 'triangle']
$ python -m pytest -q
3 passed

怎么检查结果？

• train 与 valid 类别集合必须一致（本检查器约定）。
• 每个类别至少一张非空图。

调试示例

多出来的 valid 类。 标注工具有时只在验证集加了新类，训练 vocab 与验证不一致。检查器直接对比集合，避免 fit 到一半才 KeyError/错位。若你故意用 valid_pct 不建 valid/，本脚本会提示缺少 valid——按项目选择忽略或改脚本。

常见问题

模板规则

• 训练前先 check_dataset。
• 保持 train/valid 类别一致。
• 拒绝空图。

项目资源

目录对了，训练才开始谈损失和准确率。