Fast AutoAugment (被 NeurIPS 2019 接受)

PyTorch 中官方的 Fast AutoAugment 实现。

• Fast AutoAugment 基于密度匹配，采用更高效的搜索策略来学习数据增强策略。
• 在保持相当性能的同时，Fast AutoAugment 将搜索时间提升了多个数量级。

结果

CIFAR-10 / 100

搜索：3.5 GPU 小时（比 AutoAugment 快 1428 倍），在缩减版 CIFAR-10 上使用 WResNet-40x2

模型（CIFAR-10）	基线	Cutout	AutoAugment	Fast AutoAugment (迁移/直接)
Wide-ResNet-40-2	5.3	4.1	3.7	3.6 / 3.7	下载
Wide-ResNet-28-10	3.9	3.1	2.6	2.7 / 2.7	下载
Shake-Shake(26 2x32d)	3.6	3.0	2.5	2.7 / 2.5	下载
Shake-Shake(26 2x96d)	2.9	2.6	2.0	2.0 / 2.0	下载
Shake-Shake(26 2x112d)	2.8	2.6	1.9	2.0 / 1.9	下载
PyramidNet+ShakeDrop	2.7	2.3	1.5	1.8 / 1.7	下载

模型（CIFAR-100）	基线	Cutout	AutoAugment	Fast AutoAugment (迁移/直接)
Wide-ResNet-40-2	26.0	25.2	20.7	20.7 / 20.6	下载
Wide-ResNet-28-10	18.8	18.4	17.1	17.3 / 17.3	下载
Shake-Shake(26 2x96d)	17.1	16.0	14.3	14.9 / 14.6	下载
PyramidNet+ShakeDrop	14.0	12.2	10.7	11.9 / 11.7	下载

ImageNet

搜索：450 GPU 小时（比 AutoAugment 快 33 倍），在缩减版 ImageNet 上使用 ResNet-50

模型	基线	AutoAugment	Fast AutoAugment (Top1/Top5)
ResNet-50	23.7 / 6.9	22.4 / 6.2	22.4 / 6.3	下载
ResNet-200	21.5 / 5.8	20.0 / 5.0	19.4 / 4.7	下载

注释

• 我们分别以 224 和 320 的分辨率评估了 resnet-50 和 resnet-200。根据原始 resnet 论文，resnet 200 是在 320 分辨率下测试的。此外，当我们使用该分辨率时，我们的 resnet-200 基线表现也相似。
• 然而，经过最近的代码清理和错误修复，我们发现即使使用 224x224，基线的表现仍然与之前相似。
• 当我们使用 224x224 时，resnet-200 的表现是 20.0 / 5.2。训练好的模型下载链接在此：这里。

我们还进行了使用 EfficientNet 的额外实验。

模型	基线	AutoAugment		我们的基线（批次）	+Fast AA
B0	23.2	22.7		22.96	22.68

SVHN 测试

搜索：1.5 GPU 小时

	基线	AutoAug / 我们	Fast AutoAugment
Wide-Resnet28x10	1.5	1.1	1.1

运行

我们在以下环境下进行了实验：

• python 3.6.9
• pytorch 1.2.0, torchvision 0.4.0, cuda10

搜索数据增强策略

请阅读 ray 的文档以构建合适的 ray 集群：https://github.com/ray-project/ray，并使用 master 的 redis 地址运行 search.py。

$ python search.py -c confs/wresnet40x2_cifar10_b512.yaml --dataroot ... --redis ...

使用找到的策略训练模型

您可以使用我们在 CIFAR-10 / 100 和 ImageNet 上搜索到的策略来训练网络架构。

• fa_reduced_cifar10：缩减版 CIFAR-10（4k 张图像），WResNet-40x2
• fa_reduced_imagenet：缩减版 ImageNet（50k 张图像，120 个类别），ResNet-50

$ export PYTHONPATH=$PYTHONPATH:$PWD
$ python FastAutoAugment/train.py -c confs/wresnet40x2_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar10
$ python FastAutoAugment/train.py -c confs/wresnet40x2_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar100
$ python FastAutoAugment/train.py -c confs/wresnet28x10_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar10
$ python FastAutoAugment/train.py -c confs/wresnet28x10_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar100
...
$ python FastAutoAugment/train.py -c confs/resnet50_b512.yaml --aug fa_reduced_imagenet
$ python FastAutoAugment/train.py -c confs/resnet200_b512.yaml --aug fa_reduced_imagenet

通过添加 --only-eval 和 --save 参数，您可以在不进行训练的情况下测试已训练的模型。

如果您想使用多 GPU/节点进行训练，请使用 torch.distributed.launch，例如：

$ python -m torch.distributed.launch --nproc_per_node={num_gpu_per_node} --nnodes={num_node} --master_addr={master} --master_port={master_port} --node_rank={0,1,2,...,num_node} FastAutoAugment/train.py -c confs/efficientnet_b4.yaml --aug fa_reduced_imagenet

引用

如果您在研究中使用此代码，请引用我们的论文 paper。

@inproceedings{lim2019fast,
  title={Fast AutoAugment},
  author={Lim, Sungbin and Kim, Ildoo and Kim, Taesup and Kim, Chiheon and Kim, Sungwoong},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2019}
}

问题反馈联系方式

• Ildoo Kim, ildoo.kim@kakaobrain.com

参考文献与开源项目

我们通过增大批量大小并相应调整学习率来加速训练。其余超参数则尽可能与 AutoAugment 保持一致。对于未知的超参数，我们沿用原始参考文献中的取值，或对其进行调优以达到基线性能。

• ResNet：论文1、论文2、代码
• PyramidNet：论文、代码
• Wide-ResNet：代码
• Shake-Shake：代码
• ShakeDrop 正则化：论文、代码
• AutoAugment：代码
• Ray：代码
• HyperOpt：代码

Fast AutoAugment 快速上手指南

Fast AutoAugment 是一种高效的数据增强策略搜索算法，能在保持与 AutoAugment 相当性能的同时，将搜索速度提升数个数量级。本项目基于 PyTorch 实现，适用于 CIFAR-10/100、ImageNet 等数据集。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.6.9 (兼容 3.6+)
• PyTorch: 1.2.0+
• Torchvision: 0.4.0+
• CUDA: 10.0+ (如需 GPU 加速)
• 核心依赖:
  • ray: 用于分布式搜索集群管理
  • hyperopt: 用于超参数优化

国内加速建议： 安装 Python 依赖时，推荐使用清华或阿里镜像源以提升下载速度：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

(注：若项目中无 requirements.txt，请手动安装上述核心依赖)

安装步骤

1. 克隆仓库

   git clone https://github.com/kakaobrain/fast-autoaugment.git
   cd fast-autoaugment
   

2. 配置环境变量 为了确保模块能正确导入，请将当前目录添加到 PYTHONPATH：

   export PYTHONPATH=$PYTHONPATH:$PWD
   

3. 安装依赖 确保已安装 PyTorch 和 torchvision，然后安装其他必要库（如 ray, hyperopt）：

   pip install ray hyperopt
   

基本使用

Fast AutoAugment 的工作流程分为两步：搜索增强策略 和 使用策略训练模型。

1. 搜索数据增强策略 (Search)

首先需要启动一个 Ray 集群（参考 Ray 文档），然后运行搜索脚本。以下是在 CIFAR-10 上使用 Wide-ResNet-40x2 进行搜索的示例：

python search.py -c confs/wresnet40x2_cifar10_b512.yaml --dataroot /path/to/dataset --redis <master_redis_address>

• --dataroot: 替换为您的数据集路径。
• --redis: 替换为 Ray 主节点的 Redis 地址。
• 说明: 搜索过程在 Reduced CIFAR-10 (4k 图像) 上进行，仅需约 3.5 GPU 小时。

2. 使用找到的策略训练模型 (Train)

搜索完成后，可以使用生成的策略（通过 --aug 参数指定）在完整数据集上训练模型。

CIFAR-10/100 训练示例：

# 训练 CIFAR-10
python FastAutoAugment/train.py -c confs/wresnet40x2_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar10

# 训练 CIFAR-100
python FastAutoAugment/train.py -c confs/wresnet40x2_cifar10_b512.yaml --aug fa_reduced_cifar10 --dataset cifar100

ImageNet 训练示例：

python FastAutoAugment/train.py -c confs/resnet50_b512.yaml --aug fa_reduced_imagenet

多 GPU/多节点分布式训练： 如果您需要在多卡或多机环境下训练，请使用 torch.distributed.launch：

python -m torch.distributed.launch --nproc_per_node={num_gpu_per_node} --nnodes={num_node} --master_addr={master} --master_port={master_port} --node_rank={0,1,2,...,num_node} FastAutoAugment/train.py -c confs/efficientnet_b4.yaml --aug fa_reduced_imagenet

• 请将 {...} 替换为实际的数值和地址。

仅评估已有模型： 如果您只想测试已训练好的模型而不进行训练，可添加 --only-eval 和 --save 参数。