比较 PyTorch 中的不同自动图像增强方法

中文精读

原文约 2,055 词 · 12 节

1. 开篇：为什么需要数据增强

   • 减少过拟合的最好办法之一是收集更多高质量数据，但实际中往往不可行或成本很高。
   • 数据增强是从已有数据生成新的数据记录或特征，在不额外采集数据的前提下扩充数据集。
   • 它通过构造原始输入的各种变体、让模型更难记住训练样本中的无关信息，从而改善泛化能力。
   • 数据增强在图像和文本数据中很常见，表格数据也有相应做法。
   • 本文比较 PyTorch 中的四种自动图像增强技术：AutoAugment、RandAugment、AugMix、TrivialAugment。
   • 选题缘起是作者分享了自己使用 AutoAugment 的良好体验，随后在建议下进一步考察了 torchvision 中相关的其他方法。

2. 四种自动增强方法的横向比较

   • 四种方法均已在 PyTorch/torchvision 核心库中实现，便于直接采用。
   • 在展开原理之前，先给出并排的性能对比结果。
   • 对比实验使用了一个不带额外技巧的简单 ResNet-18。
   • 为保持简单，实验中没有加入其他减少过拟合的手段。

3. 代码

   • 作者称实验代码很简单，已放在 GitHub 上，供读者运行、改写或修改；正文只展示并概述主要部分，完整可执行代码在 GitHub。
   • 训练变换先用 Compose 把图像缩放到 32，再在“不加增强”与 AutoAugment、RandAugment、AugMix、TrivialAugmentWide 之一之间选择，最后转成张量；验证与测试变换只做缩放到 32 和转张量。
   • 用 Lightning 的 seed_everything(123) 设置种子；Cifar10DataModule 的批大小为 256、num_workers 为 4。
   • 模型用 torchvision 的无预训练权重 ResNet-18，把最后一层全连接换成 512→10 的输出层（对应 CIFAR-10 的 10 类），并用学习率 0.05 封装成 LightningModel。
   • Trainer 设置为最多 1000 轮、使用 accelerator="gpu" 且 devices=[0]（单 GPU）、日志写入 CSVLogger、deterministic=True；随后调用 trainer.fit 训练，并用 trainer.test 做模型评估。
   • 四个训练场景共用同一组超参数；作者不逐行讲解代码，疑问可通过配套的 GitHub Discussion 论坛提出。
   python

   import torch, torchvision from torchvision import transforms import lightning as L ## Dataset train_transform = torchvision. transforms. Compose([transforms. Resize(32),# One of the following: # 1) Nothing # 2) transforms.AutoAugment() # 3) transforms.RandAugment() # 4) transforms.AugMix() # 5) transforms.TrivialAugmentWide() transforms. ToTensor(),]) valid_and_test_transform = transforms. Compose([transforms. Resize(32), transforms. ToTensor(),]) L. seed_everything(123) dm = Cifar10DataModule(batch_size = 256, train_transform = train_transform, valid_and_test_transform = valid_and_test_transform, num_workers = 4) ## Model pytorch_model = torchvision. models. resnet18(weights = None) pytorch_model. fc = torch. nn. Linear(512, 10) lightning_model = LightningModel(model = pytorch_model, learning_rate =0.05) ## Training trainer = L. Trainer(max_epochs = 1000, accelerator = "gpu", devices =[0], logger = CSVLogger(save_dir ="logs/"), deterministic = True,) trainer. fit(model = lightning_model, datamodule = dm) ## Model evaluation trainer. test(model = lightning_model, datamodule = dm)

4. 实验结果

   • AutoAugment 相对不做增强把测试集准确率提高约 12%，提升幅度很大。
   • 作为 AutoAugment 后继方法的 RandAugment 又提升约 2 个百分点；更新且简单得多的 TrivialAugment 再比 RandAugment 好约 1 个百分点。
   • AugMix 优于不做增强，但不如 AutoAugment、RandAugment、TrivialAugment；这在预期之内，因为 AugMix 是针对分布偏移而非提升验证准确率设计的。
   • 作为对照基线还加入了 RandomAffine（对图像做包含随机平移、缩放与剪切/shear 的随机仿射变换），它优于不增强，但远不及 AutoAugment、RandAugment、TrivialAugment。
   • 验证准确率曲线仍有上升趋势，因此又把 AutoAugment、RandAugment、AugMix、TrivialAugment 跑了 2000 轮（结果未展示），但测试集表现相比 1000 轮并未提升。

5. 局限性

   • 实验只用了单一网络结构（ResNet-18）和单一数据集（CIFAR-10）。
   • 换成不同模型结构、数据集或配合其他减少过拟合的技术时，结果可能不同。

6. 参考文献

   • AutoAugment：Cubuk 等（2019 年 4 月），实现为 torchvision.transforms.AutoAugment。
   • RandAugment：Cubuk 等（2019 年 11 月），实现为 torchvision.transforms.RandAugment。
   • AugMix：Hendrycks 等（2020），实现为 torchvision.transforms.AugMix。
   • TrivialAugment：Mueller、Hutter（2021 年 8 月），实现为 torchvision.transforms.TrivialAugmentWide。

7. 方法小结与比较

   • 本节依据上述文献，简述各自动或学习式图像增强方法的做法，并相互比较。

8. AutoAugment

   • 原始论文中，作者学习能优化给定数据集验证准确率的图像变换组合。
   • 候选变换包括剪切、平移、旋转、对比度、像素翻转、直方图均衡、曝光过度、色调分离、颜色、亮度、锐度、Cutout 以及 Sample Pairing。
   • 寻找有效变换组合是一个离散优化问题，可用网格搜索、随机搜索、强化学习或进化算法求解。
   • AutoAugment 选用强化学习，以循环神经网络作为控制器；由于准确率不可微，用策略梯度方法更新 RNN。
   • 控制器是含 100 个隐藏单元的 1 层 LSTM，训练使用 PPO（近端策略优化）。
   • 策略包含两步：施加某种增强的概率，以及所选增强的强度；论文针对 CIFAR-10/100、SVHN 和 ImageNet 学习策略，并称这些策略可迁移到新数据集。

9. AutoAugment（续）

   • 在 PyTorch/torchvision 中调用 AutoAugment 时使用的是对应数据集已学好的策略，而不是现场训练；例如 CIFAR-10 用 AutoAugmentPolicy.CIFAR10。

10. RandAugment

    • AutoAugment 通过强化学习的策略优化在代理任务上学习增强策略。
    • RandAugment 由 AutoAugment 的作者提出，思路简单得多：把增强技术的选择转化为超参数搜索问题。
    • 因此它不需要单独搜索；对每个小批量中的每张图像，以均匀概率选取一组图像变换，每种增强被选中的机会相同。
    • 它有两个超参数：对一张图像施加的增强数量，以及这些增强的强度。
    • 论文称该方法可用两行 Python 代码实现；每张变换的强度被离散成 31 个档位，默认增强数量为 2、默认强度档位为 9。
    python

    transforms = ['Identity', 'AutoContrast', 'Equalize', 'Rotate', 'Solarize', 'Color', 
                  'Posterize', 'Contrast', 'Brightness', 'Sharpness', 'ShearX', 'ShearY', 
                  'TranslateX', 'TranslateY']
    
    def randaugment(N, M):
        """Generate a set of distortions.
    
        Args:
          N: Number of augmentation transformations to apply sequentially.
          M: Magnitude for all the transformations.
        """
        sampled_ops = np.random.choice(transforms, N)
        return [(op, M) for op in sampled_ops]

11. AugMix

    • AugMix 专门用于提升对数据偏移（例如部署中常见的偏移）的鲁棒性。
    • 它把不同的图像增强技术混合起来以改善分类器表现，同时优化一致性损失。
    • 所用变换与 AutoAugment 类似，但去掉了对比度、颜色、亮度、锐度和 Cutout；且不是单纯串行施加，而是把增强结果组合起来。
    • 最终图像由增强图像与原图插值（混合）得到，而增强图像本身也以不同的混合权重组合多个分支。
    • 总体损失＝分类损失＋一致性损失：分类损失是标准交叉熵，一致性损失是模型对原图与对增强图输出的 softmax 之间的 Jensen-Shannon 散度。
    • Jensen-Shannon 散度是衡量两个概率分布差异的 KL 散度的对称版本，而 KL 散度等于交叉熵减去自熵项。

12. TrivialAugment

    • TrivialAugment 是极简做法，每张图像只施加一次增强，且不是学习式过程。
    • 表面看它像是 RandAugment 的特例。
    • 区别一：RandAugment 可能对一张图像施加多个变换，而 TrivialAugment 严格只用一个。
    • 区别二：RandAugment 对整个数据集使用固定的增强强度（作为超参数），TrivialAugment 则为每张图像随机采样强度。