基于种群的增强（PBA）

新增：使用笔记本 pba.ipynb 可视化 PBA 及其应用的增强策略！

现已支持 Python 3。

目录

1. 简介
2. 快速入门
3. 复现结果
4. 运行 PBA 搜索
5. 引用

简介

基于种群的增强（PBA）是一种能够快速高效地为神经网络训练学习数据增强函数的算法。PBA 在计算资源消耗减少一千倍的情况下，仍能取得与 CIFAR 数据集上最先进方法相当的性能，使研究人员和从业者仅需一台工作站的 GPU 即可有效地学习新的增强策略。

本仓库包含论文《基于种群的增强：高效学习增强策略调度》（http://arxiv.org/abs/1905.05393）在 TensorFlow 和 Python 中的实现代码。其中既包括使用文中报告的增强策略进行模型训练，也包括发现新的增强策略调度。

下方展示了我们的增强策略可视化效果。

快速入门

代码同时支持 Python 2 和 3。

安装依赖

pip install -r requirements.txt

下载 CIFAR-10/CIFAR-100 数据集

bash datasets/cifar10.sh
bash datasets/cifar100.sh

复现结果

用于复现结果的脚本位于 scripts/table_*.sh 中。所有脚本均需一个参数——模型名称。可用选项为论文表 2 中针对各数据集报告的模型，包括：wrn_28_10, ss_32, ss_96, ss_112, pyramid_net。超参数也内置于每个脚本文件中。

例如，要复现 CIFAR-10 上 Wide-ResNet-28-10 的结果：

bash scripts/table_1_cifar10.sh wrn_28_10

要复现缩减版 SVHN 上 Shake-Shake (26 2x96d) 的结果：

bash scripts/table_4_svhn.sh rsvhn_ss_96

建议从缩减版 SVHN 上 Wide-ResNet-28-10 开始，这在 Titan XP GPU 上不到 10 分钟即可完成，测试准确率可达 91% 以上。

运行规模较大的模型时，若采用 1800 轮训练，则可能需要数天时间。例如，CIFAR-10 上 PyramidNet+ShakeDrop 在 Tesla V100 GPU 上大约需要 9 天。

运行 PBA 搜索

使用 scripts/search.sh 文件对 Wide-ResNet-40-2 进行 PBA 搜索。需提供一个参数——数据集名称，可选值为 rsvhn 或 rcifar10。

通过指定部分 GPU 资源，可在同一 GPU 上启动多个试验。缩减版 SVHN 大约需要一小时在 Titan XP GPU 上完成，而缩减版 CIFAR-10 则需约 5 小时。

CUDA_VISIBLE_DEVICES=0 bash scripts/search.sh rsvhn

搜索过程中使用的最终调度可从 Ray 的结果目录中获取，日志文件则可通过 pba/utils.py 中的 parse_log() 函数转换为增强策略调度。例如，在缩减版 CIFAR-10 上经过 200 轮训练学到的增强策略调度，被拆分为概率和幅度两个超参数，并分别可视化如下：

随时间变化的概率超参数	随时间变化的幅度超参数

引用

如果您在研究中使用了 PBA，请引用以下文献：

@inproceedings{ho2019pba,
  title     = {Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules},
  author    = {Daniel Ho and
               Eric Liang and
               Ion Stoica and
               Pieter Abbeel and
               Xi Chen
  },
  booktitle = {ICML},
  year      = {2019}
}

PBA (Population Based Augmentation) 快速上手指南

PBA 是一种高效的数据增强策略学习算法，能够在单张 GPU 上以极低的计算成本（相比传统方法减少 1000 倍算力）学习到最优的数据增强调度方案，并在 CIFAR 等数据集上达到业界领先的准确率。

环境准备

• 操作系统: Linux / macOS (Windows 需通过 WSL 或 Docker 运行)
• Python 版本: 支持 Python 2 和 Python 3 (推荐使用 Python 3)
• 深度学习框架: TensorFlow
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (推荐 Titan XP, Tesla V100 或更高)
• 依赖管理: pip

提示: 国内用户建议在安装依赖时指定清华或阿里镜像源以加速下载。

安装步骤

1. 克隆代码库 (假设已获取源码目录)

   cd pba
   

2. 安装依赖包 建议使用国内镜像源加速安装：

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

3. 下载数据集 执行脚本自动下载 CIFAR-10 和 CIFAR-100 数据集：

   bash datasets/cifar10.sh
   bash datasets/cifar100.sh
   

   如果下载速度慢，可手动下载数据集后放入 datasets/ 对应目录，或修改脚本中的下载链接为国内镜像地址。

基本使用

1. 复现经典结果 (Reproduce Results)

项目提供了预定义的脚本来复现论文中的实验结果。所有脚本位于 scripts/ 目录下，运行时需传入模型名称参数。

可用模型参数:

• wrn_28_10 (Wide-ResNet-28-10)
• ss_32, ss_96, ss_112 (Shake-Shake 不同配置)
• pyramid_net (PyramidNet+ShakeDrop)

示例 A: 在 CIFAR-10 上复现 Wide-ResNet-28-10 的结果

bash scripts/table_1_cifar10.sh wrn_28_10

示例 B: 在 Reduced SVHN 上复现 Shake-Shake (26 2x96d) 的结果 这是一个适合快速测试的示例，在 Titan XP GPU 上约 10 分钟即可完成，测试准确率可达 91% 以上。

bash scripts/table_4_svhn.sh rsvhn_ss_96

注意: 大型模型（如 CIFAR-10 PyramidNet）可能需要数天的训练时间。

2. 运行 PBA 搜索 (Run PBA Search)

使用 PBA 算法自动搜索特定数据集上的最优增强策略。

支持的数据集参数:

• rsvhn (Reduced SVHN)
• rcifar10 (Reduced CIFAR-10)

执行搜索: 以下命令将在 GPU 0 上启动针对 Reduced SVHN 的搜索任务（约需 1 小时）：

CUDA_VISIBLE_DEVICES=0 bash scripts/search.sh rsvhn

查看结果: 搜索完成后，策略调度文件将保存在 Ray 的结果目录中。你可以使用 pba/utils.py 中的 parse_log() 函数解析日志，将其转换为可视化的概率和幅度超参数曲线。

3. 可视化分析 (可选)

项目包含一个 Jupyter Notebook (pba.ipynb)，可用于可视化 PBA 搜索过程及最终应用的数据增强效果。建议在 Jupyter Lab 或 Notebook 中打开查看。