第4章数据加载与预处理

第4章：数据加载与预处理——喂给模型什么，它就学会什么

做量化感知训练，很多人一上来就调参数、改代码，结果模型训出来精度掉得一塌糊涂。我见过太多这种情况了。其实啊，问题往往出在数据上——你喂给模型的数据不对，后面再怎么折腾也是白搭。

这一章，咱们就聊聊数据加载与预处理。说白了，就是搞清楚三件事：用什么数据、怎么增强、校准集怎么搞。

4.1 标准数据集加载：ImageNet 与 CIFAR-10

做 QAT，最常用的两个数据集就是 ImageNet 和 CIFAR-10。ImageNet 是大哥，1000 类、百万级图片，适合做大规模分类。CIFAR-10 是小弟，10 类、6 万张 32x32 的小图，适合快速验证。

我个人习惯，在项目初期先用 CIFAR-10 跑通流程，再换 ImageNet 做最终验证。为什么呢？因为 CIFAR-10 训练快，几分钟就能看到效果，调试起来效率高。

4.1.1 用 PyTorch 加载 CIFAR-10

PyTorch 的 torchvision 直接提供了 CIFAR-10 的接口，几行代码就能搞定：

import torch
import torchvision
import torchvision.transforms as transforms

# 定义数据预处理
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

# 加载训练集和测试集
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform_train)
testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform_test)

trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=128, shuffle=True, num_workers=4)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=100, shuffle=False, num_workers=4)

这里要注意，num_workers 我一般设成 4 或 8，具体看你的 CPU 核心数。设太少加载慢，设太多反而会出问题。我曾经在 32 核的机器上设了 32，结果内存爆了……嗯，教训啊。

4.1.2 加载 ImageNet

ImageNet 没有自动下载接口，你得自己去官网下载。下载后按文件夹组织好：

# 假设你的 ImageNet 数据在 /data/imagenet 下
# 结构应该是：
# /data/imagenet/train/n01440764/n01440764_10026.JPEG
# /data/imagenet/val/n01440764/ILSVRC2012_val_00000293.JPEG

traindir = '/data/imagenet/train'
valdir = '/data/imagenet/val'

# 训练集增强
transform_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

# 验证集只做 resize 和 center crop
transform_val = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

train_dataset = torchvision.datasets.ImageFolder(traindir, transform_train)
val_dataset = torchvision.datasets.ImageFolder(valdir, transform_val)

注意：ImageNet 的验证集默认没有按子文件夹分类，你需要用官方脚本整理一下。不然 ImageFolder 会报错。我当初第一次跑的时候就被这个坑了半小时。

4.2 数据增强策略：不只是为了涨点

数据增强在 QAT 里有两个作用：一是提升模型泛化能力，二是让量化感知训练更鲁棒。你想想看，如果训练时模型只见过「标准」图片，量化后遇到一点噪声就崩了，那还怎么部署？

4.2.1 基础增强方法

方法	说明	QAT 中的建议
RandomCrop	随机裁剪图片	必须用，模拟不同视角
RandomHorizontalFlip	随机水平翻转	必须用，简单有效
ColorJitter	调整亮度、对比度、饱和度	建议用，增强对光照变化的鲁棒性
RandomRotation	随机旋转	小角度可用，大角度会破坏语义
RandomErasing	随机遮挡部分区域	对量化有帮助，模拟信息丢失

我个人经验是，不要一股脑把所有增强都用上。比如对 CIFAR-10，RandomCrop + RandomHorizontalFlip 就够用了。对 ImageNet，再加个 ColorJitter 效果更好。加太多反而可能让模型学不到关键特征。

4.2.2 高级增强：AutoAugment 与 RandAugment

如果你想让模型更「抗造」，可以试试 AutoAugment 或 RandAugment。这些方法会自动搜索最优的增强策略。

from torchvision.transforms import autoaugment

# 使用 RandAugment
transform_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    autoaugment.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

我在一个移动端部署项目里试过 RandAugment，量化后的模型精度比只用基础增强高了 0.8%。别小看这 0.8%，在部署场景里可能就是「能用」和「不能用」的区别。

小技巧：QAT 训练时，增强强度可以比普通训练稍微弱一点。因为量化本身会引入噪声，增强太强反而让模型难以收敛。我一般把 RandAugment 的 magnitude 从 9 降到 7。

4.3 校准数据集：QAT 的灵魂

校准数据集（Calibration Dataset）是 QAT 里最容易被忽视的概念。很多人以为随便拿一批数据就能校准，结果量化后精度掉成狗。

说白了，校准数据集就是用来统计激活值分布的那批数据。QAT 在训练过程中需要知道每一层输出的数值范围，才能确定量化参数（scale 和 zero_point）。

4.3.1 校准数据集的要求

代表性：必须能代表真实部署场景的数据分布。比如你做的是人脸识别，校准集里就不能全是风景照。
多样性：覆盖所有类别，每个类别至少 10-20 张图。我见过有人只用 100 张图校准 1000 类的模型，结果可想而知。
数量适中：一般 500-2000 张就够了。太多反而增加校准时间，太少统计不准。

4.3.2 如何构建校准数据集

最直接的方法是从训练集中随机抽取一个子集：

import random
from torch.utils.data import Subset

# 假设 trainset 是完整的训练集
total_size = len(trainset)
calibration_size = 1000  # 校准集大小

# 随机抽取索引
indices = random.sample(range(total_size), calibration_size)
calibration_dataset = Subset(trainset, indices)

# 创建 DataLoader
calibration_loader = torch.utils.data.DataLoader(
    calibration_dataset, batch_size=64, shuffle=False, num_workers=4)

但这里有个坑：随机抽取可能漏掉某些类别。尤其是类别不平衡的数据集，小类别可能一个都没抽到。

我建议用分层采样：

from torch.utils.data import Subset, DataLoader
import numpy as np

def create_balanced_calibration_set(dataset, samples_per_class=10):
    """
    为每个类别均匀采样，构建校准集
    """
    # 获取所有标签
    labels = [dataset[i][1] for i in range(len(dataset))]
    labels = np.array(labels)
    
    # 获取类别列表
    classes = np.unique(labels)
    
    calibration_indices = []
    for cls in classes:
        # 找到该类别的所有索引
        cls_indices = np.where(labels == cls)[0]
        # 随机选择 samples_per_class 个
        selected = np.random.choice(cls_indices, 
                                    size=min(samples_per_class, len(cls_indices)),
                                    replace=False)
        calibration_indices.extend(selected.tolist())
    
    return Subset(dataset, calibration_indices)

# 使用：每个类别取 10 张
calibration_dataset = create_balanced_calibration_set(trainset, samples_per_class=10)
calibration_loader = DataLoader(calibration_dataset, batch_size=64, shuffle=False)

核心要点：校准集的质量直接决定 QAT 的成败。宁可用 500 张精心挑选的图片，也不用 2000 张随便凑的图片。我在一个工业检测项目里，就因为校准集里少了「缺陷样本」这个类别，量化后模型对缺陷的检测率直接掉了 15%。后来补上这个类别的数据，精度就回来了。

4.3.3 校准数据集的预处理

校准集的数据增强要特别注意：不要用随机增强。因为校准的目的是统计激活值的真实分布，随机增强会引入额外的方差，导致统计不准。

# 校准集只做固定预处理
transform_calibration = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

嗯，这里要特别强调：校准集和验证集可以用相同的预处理，但绝对不能和训练集用相同的随机增强。

4.4 本章知识体系

下面这张图总结了数据加载与预处理的核心逻辑，我画了个流程图，方便你理解各个环节的关系：

从这张图可以看得很清楚：训练集用随机增强，校准集和验证集用固定预处理。这个原则一定要记住，我在好几个项目里看到有人把校准集也做了随机增强，结果量化参数飘忽不定，精度忽高忽低。

避坑指南：我曾经在一个自动驾驶项目里，校准集用了 5000 张白天场景的图片，结果模型在夜间场景下量化后精度暴跌。后来我重新构建了校准集，白天和夜间各 1000 张，问题就解决了。记住：校准集要覆盖所有部署场景。

好了，数据加载与预处理就聊到这儿。下一章咱们会深入 QAT 的核心——如何把量化操作嵌入到训练过程中。到时候你会看到，数据准备得越好，后面的工作就越顺利。

📚 课程目录