内存优化型批次归一化激活

内存优化型深度神经网络训练中的内嵌式批次归一化激活

内嵌式批次归一化激活（InPlace-ABN）是一种新颖的方法，用于减少训练深层网络所需的内存。它通过重新定义BN（批次归一化）加非线性激活作为单一的就地操作，在现代架构如ResNet、ResNeXt和Wider ResNet中实现高达50%的内存节省，并根据需要智能地丢弃或重算中间缓冲区。

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

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

我们现在还发布了语义分割的推断代码，以及在Mapillary Vistas上训练的模型，该模型在Mapillary Vistas语义分割排行榜上取得第1名的成绩。更多信息可在页面底部找到。

引用

如果你在研究中使用了In-Place Activated 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-激活-卷积序列时，大多数深度学习框架需要存储两个大的缓冲区，即BN的输入x和卷积的输入z。这是因为在BN和卷积的标准反向传播实现中依赖于它们的输入来计算梯度。使用Inplace-ABN替换BN-激活序列，我们可以安全地丢弃x，从而在训练时节省最多50%的GPU内存。为了实现这一点，我们将BN的反向传播重写为基于其输出y的形式，而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脚本所需的一些额外库。

强制使用CUDA编译

在没有GPU访问权限的系统（例如Docker容器）上强制编译原生CUDA函数，需要设置以下两个环境变量：

export TORCH_CUDA_ARCH_LIST="{archs}"
export IABN_FORCE_CUDA=1

其中{archs}是目标CUDA架构的列表，例如Pascal;Volta，6.0;6.5等等。

在ImageNet-1k上训练

以下是我们在arXiv论文中的结果（ Top-1 / Top-5 分数），以及相应的已训练模型和MD5校验和。提供的模型文件受ImageNet附带的许可证约束。

网络	批次大小	224像素	224像素，10个裁剪	320像素	训练模型及MD5校验和
ResNeXt101，标准BN	256	77.04 / 93.50	78.72 / 94.47	77.92 / 94.28	448438885986d14db5e870b95f814f91
ResNeXt101，InPlace-ABN	512	78.08 / 93.79	79.52 / 94.66	79.38 / 94.67	3b7a221cbc076410eb12c8dd361b7e4e
ResNeXt152，InPlace-ABN	256	78.28 / 94.04	79.73 / 94.82	79.56 / 94.67	2c8d572587961ed74611d534c5b2e9ce
WideResNet38，InPlace-ABN	256	79.72 / 94.78	81.03 / 95.43	80.69 / 95.27	1c085ab70b789cc1d6c1594f7a761007
ResNeXt101，同步InPlace-ABN	256	77.70 / 93.78	79.18 / 94.60	78.98 / 94.56	0a85a21847b15e5a242e17bf3b753849
DenseNet264，InPlace-ABN	256	78.57 / 94.17	79.72 / 94.93	79.49 / 94.89	0b413d67b725619441d0646d663865bf
ResNet50v1，同步InPlace-ABN	512	75.53 / 92.59	77.04 / 93.57	76.60 / 93.49	2522ca639f7fdfd7c0089ba1f5f6c2e8
ResNet34v1，同步InPlace-ABN	512	73.27 / 91.34	75.19 / 92.66	74.87 / 92.42	61515c1484911c3cc753d405131e1dda
ResNet101v1，同步InPlace-ABN	512	77.07 / 93.45	78.58 / 94.40	78.25 / 94.19	1552ae0f3d610108df702135f56bd27b

数据准备

我们的脚本使用torchvision.datasets.ImageFolder加载ImageNet数据，需要按照以下结构组织文件夹：

根目录/训练/[类别ID1]/xxx.{jpg,png,jpeg}
根目录/训练/[类别ID1]/xxy.{jpg,png,jpeg}
根目录/训练/[类别ID2]/xxz.{jpg,png,jpeg}
...

根目录/验证/[类别ID1]/asdas.{jpg,png,jpeg}
根目录/验证/[类别ID1]/123456.{jpg,png,jpeg}
根目录/验证/[类别ID2]/__32_.{jpg,png,jpeg}
...

图像可以具有任意名称，只要扩展名是可识别的图像格式即可。类别的ID也是自由形式的，但要求在训练和验证数据之间匹配。标准ImageNet分布中的训练数据已经按要求的格式给出，而验证图像则需要按照上述方式拆分为类子文件夹。

训练

主要的训练脚本是scripts/train_imagenet.py：它支持对ImageNet或其他以上述方式格式化的数据集进行训练，并记录相关指标的Tensorboard日志，以及定期保存快照。 大多数训练参数可以通过json格式的配置文件指定（查看scripts/imagenet/config.py以获取完整的可配置参数列表）。 配置文件中未明确指定的所有参数都将设置为其默认值，也可以在scripts/imagenet/config.py中找到。

我们在arXiv上的结果可以通过运行scripts/train_imagenet.py并使用scripts/experiments中的配置文件来复现。例如，使用InPlace-ABN、Leaky ReLU和batch_size = 512训练ResNeXt101的命令如下：

cd scripts
python -m torch.distributed.launch --nproc_per_node <每节点GPU数量> 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作物验证和兼容网络之间的权重迁移（例如，从带有ReLU的ResNeXt101到带有Leaky ReLU的ResNeXt101）。 此脚本接受与scripts/train_imagenet.py相同的配置文件，但请注意，scale_val和crop_val参数会被命令行参数--scale和--crop所忽略。

例如，要使用从缩放到256像素的图像中提取的大小为224的10个作物验证上面训练的ResNeXt101，可以运行：

cd scripts
python -m torch.distributed.launch --nproc_per_node <每节点GPU数量> 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应用于了在上述宽残差网38模型之上训练的DeepLab3分割头。由于InPlace-ABN，我们可以显著增加输入到这个模型的数据量，最终在城市景观、Mapillary Vistas、AutoNUE、Kitti和ScanNet的分割排行榜上取得了第一名的位置。训练设置主要遵循我们在论文中描述的方式。

在Mapillary Vistas上预训练的模型

我们发布了一个在Mapillary Vistas研究集上训练的宽残差网38 + DeepLab3分割模型。这是我们在MVD语义分割排行榜上取得第一名所使用的模型。下面提供的分割模型文件根据CC BY-NC-SA 4.0许可协议提供。

网络	mIOU	训练模型 (+md5)
宽残差网38 + DeepLab3	53.42	913f78486a34aa1577a7cd295e8a33bb

要使用此模型，请下载上面链接的.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库，以便于在其他项目中使用
• ImageNet / Vistas脚本仍保留在scripts文件夹中
• 现在需要PyTorch 1.1

更新2019年1月8日：

• 启用了多进程处理和多进程中Inplace ABN同步（之前使用线程）。现在需要使用DistributedDataParallel而不是DataParallel
• 增加了对fp16的兼容性（目前允许fp16输入，但要求模块保持在fp32模式）

2019年2月更新：

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

我们修改了ImageNet的训练代码和BN同步，使其可以在多个进程中工作。我们也使我们的Inplace ABN模块与fp16兼容。