就地激活的 BatchNorm

用于 DNN 的内存优化训练的就地激活 BatchNorm

就地激活 BatchNorm (InPlace-ABN) 是一种减少训练深度网络所需内存的新方法。通过将 BN + 非线性激活重新定义为单个就地操作，同时根据需要智能地删除或重新计算中间缓冲区，它可以在 ResNet、ResNeXt 和 Wider ResNet 等现代架构中节省高达 50% 的内存。

该存储库包含 InPlace-ABN 层的PyTorch实现，以及一些用于重现我们论文中报告的 ImageNet 分类结果的训练脚本。

• 概述
• 安装
• 在 ImageNet 上训练

我们现在还发布了语义分割的推理代码，以及经过 Mapillary Vistas 训练的模型，在 Mapillary Vistas 语义分割排行榜上排名第一。更多信息可以在本页底部找到。

引文

如果您在研究中使用就地激活的 BatchNorm，请引用：

@inproceedings{rotabulo2017place,
  title={In-Place Activated BatchNorm for Memory-Optimized Training of DNNs},
  author={Rota Bul\`o, Samuel and Porzi, Lorenzo and Kontschieder, Peter},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  year={2018}
}

概述

在前向传递中处理一个 BN-Activation-Convolution 序列时，大多数深度学习框架需要存储两个大缓冲区，即xBN 的输入和zConv 的输入。这是必要的，因为 BN 和 Conv 的反向传播的标准实现依赖于它们的输入来计算梯度。使用 Inplace-ABN 替换 BN-Activation 序列，我们可以安全地丢弃x，从而在训练时节省高达 50% 的 GPU 内存。为了实现这一点，我们根据输出 重写了 BN 的后向传递，而输出又是 通过反转激活函数y来重建的。z

与标准 BN 相比，BN 比例因子的参数化发生了变化，以确保可逆变换。具体来说，比例因子变为 。

要求

要安装 PyTorch，请参考https://github.com/pytorch/pytorch#installation。

注 1：我们的代码需要PyTorch v1.1 或更高版本

注意 2：我们只能提供对 Linux 平台和 CUDA 版本 >= 10.0 的支持

注 3：一般来说，由于 BN 缩放参数的处理方式不同，不可能将使用标准 BN 训练的网络中的权重加载到 InPlace-ABN 网络中而不会导致性能严重下降

要安装包含 iABN 层的软件包：

pip install inplace_abn

请注意，InPlace-ABN 的某些部分具有本地 C++/CUDA 实现，这意味着上面的命令将需要编译它们。

或者，要下载并安装我们库的最新版本，还需要获取 Imagenet / Vistas 脚本的副本：

git clone https://github.com/mapillary/inplace_abn.git
cd inplace_abn
python setup.py install
cd scripts
pip install -r requirements.txt

上面的最后一个命令将安装 Imagenet / Vistas 脚本所需的一些附加库。

ImageNet-1k 训练

在这里，您可以找到我们的 arXiv 论文（top-1 / top-5 分数）的结果，分别带有相应的训练模型和 md5 校验和。下面提供的模型文件根据ImageNet 附带的许可提供。

数据准备

我们的脚本使用torchvision.datasets.ImageFolder 来加载 ImageNet 数据，它需要按如下方式组织的文件夹：

root/train/[class_id1]/xxx.{jpg,png,jpeg}
root/train/[class_id1]/xxy.{jpg,png,jpeg}
root/train/[class_id2]/xxz.{jpg,png,jpeg}
...

root/val/[class_id1]/asdas.{jpg,png,jpeg}
root/val/[class_id1]/123456.{jpg,png,jpeg}
root/val/[class_id2]/__32_.{jpg,png,jpeg}
...

图像可以有任何名称，只要扩展名是可识别的图像格式。类 id 也是自由格式的，但它们应该在训练数据和验证数据之间匹配。请注意，标准 ImageNet 分布中的训练数据已经以所需的格式给出，而验证图像需要如上所述拆分到类子文件夹中。

训练

主要的训练脚本是scripts/train_imagenet.py：这支持在 ImageNet 或任何其他如上所述格式化的数据集上进行训练，同时以 Tensorboard 格式保存相关指标的日志并定期保存快照。大多数训练参数可以指定为json-formatted 配置文件（ 在此处查看可配置参数的完整列表）。所有未在配置文件中明确指定的参数都设置为其默认值，也可在 scripts/imagenet/config.py中找到。

scripts/train_imagenet.py我们的 arXiv 结果可以通过运行scripts/experiments. 例如，ResNeXt101使用 InPlace-ABN、Leaky ReLU进行训练的命令batch_size = 512是：

cd scripts
python -m torch.distributed.launch --nproc_per_node <n. GPUs per node> train_imagenet.py --log-dir /path/to/tensorboard/logs experiments/resnext101_ipabn_lr_512.json /path/to/imagenet/root

验证

验证scripts/train_imagenet.py在每个训练时期结束时运行。要验证经过训练的模型，您可以使用该scripts/test_imagenet.py脚本，该脚本允许 10-crops 验证和跨兼容网络传输权重（例如，从ResNeXt101ReLU 到ResNeXt101Leaky ReLU）。此脚本接受与 相同的配置文件scripts/train_imagenet.py，但请注意scale_val和 crop_val参数被忽略而有利于--scale和--crop命令行参数。

例如，要使用缩放到像素的图像ResNeXt101的 10 个大小来验证上面的训练 ，您可以运行：224256

cd scripts
python -m torch.distributed.launch --nproc_per_node <n. GPUs per node> test_imagenet.py --crop 224 --scale 256 --ten_crops experiments/resnext101_ipabn_lr_512.json /path/to/checkpoint /path/to/imagenet/root

用于城市景观和 Mapillary Vistas 的语义分割

我们已经成功地将 InPlace-ABN 与在上述 WideResNet38 模型之上训练的 DeepLab3 分割头一起使用。由于 InPlace-ABN，我们可以显着增加该模型的输入数据量，最终使我们能够在Cityscapes、 Mapillary Vistas、AutoNUE、 Kitti和 ScanNet分割排行榜上获得第一名。训练设置大多遵循我们论文中的描述。

Mapillary Vistas 预训练模型

我们发布了在 Mapillary Vistas 研究集上训练的 WideResNet38 + DeepLab3 分割模型。这是用于在 MVD 语义分割排行榜上达到 #1 位置的模型。下面提供的分段模型文件在 CC BY-NC-SA 4.0 许可下提供。

要使用它，请下载.pth.tar上面链接的模型文件并运行test_vistas.py脚本如下：

cd scripts
python test_vistas.py /path/to/model.pth.tar /path/to/input/folder /path/to/output/folder

该脚本将处理输入文件夹中的所有.png,.jpg和.jpeg图像，并将输出文件夹中的预测作为.png图像写入。有关其他选项，例如测试时间增加，请查阅脚本的帮助消息。

上面写的测试数据的结果是通过仅使用比例 1.0 + 翻转获得的。

变更日志

2019 年 7 月 4 日更新：1.0.0 版

• 在 Pytorch 最新的本机 BN 实现之后完全重写 CUDA 代码
• 改进的同步 BN 实现，正确处理不同的每 GPU 批量大小和 Pytorch 分布式组
• iABN 层现在打包在一个可安装的 python 库中，以简化在其他项目中的使用
• scriptsImagenet / Vistas 脚本在文件夹中仍然可用
• 现在需要 PyTorch 1.1

2019 年 1 月 8 日更新：

• 在多个进程上启用多处理并就地 ABN 同步（以前使用线程）。它现在需要使用 DistributedDataParallel 而不是 DataParallel
• 增加了对 fp16 的兼容性（目前允许 fp16 输入但需要模块保持 fp32 模式）
• 现在需要 PyTorch 1.0

2019 年 2 月更新：

• 添加了 ResNet34v1、ResNet50v1 和 ResNet101v1 ImageNet-1k 预训练模型

我们修改了 imagenet 训练代码和 BN 同步，以便与多个进程一起工作。我们还增加了 Inplace ABN 模块与 fp16 的兼容性。