PyTorch 79. NVIDIA训练深度学习模型神器APEX详细介绍

APEX 是来自英伟达 (NVIDIA) 的一个很好用的深度学习加速库。由英伟达开源，完美支持PyTorch框架，用于改变数据格式来减小模型显存占用的工具。其中最有价值的是 amp (Automatic Mixed Precision) ，将模型的大部分操作都用 Float16 数据类型测试，一些特别操作仍然使用 Float32。并且用户仅仅通过三行代码即可完美将自己的训练代码迁移到该模型。实验证明，使用 Float16 作为大部分操作的数据类型，并没有降低参数，在一些实验中，反而由于可以增大 Batch size，带来精度上的提升，以及训练速度上的提升。

Apex Github地址：

Apex 英文介绍：

前提是你安装好了CUDA和CUDNN，以及你的系统是Ubuntu系统。

Apex 还通过以下方式支持仅 Python 构建 (Pytorch 0.4 需要)。

安装之后，clone下来的apex文件夹就可以删除了。

查看能否正确导入apex：

1. Amp: Automatic Mixed Precision

apex.amp 是一种通过仅更改脚本的 3 行来启用混合精度训练的工具。 通过向 amp.initialize 提供不同的 flags，用户可以轻松地试验不同的纯精度和混合精度训练模式。

API 文档：

2. Distributed Training

apex.parallel.DistributedDataParallel 是一个模块包装器，类似于 torch.nn.parallel.DistributedDataParallel。 它支持方便的多进程分布式训练，针对 NVIDIA 的 NCCL 通信库进行了优化。

API 文档：

apex.parallel.DistributedDataParallel 是一个模块包装器，可实现轻松的多进程分布式数据并行训练，类似于 torch.nn.parallel.DistributedDataParallel。 参数在初始化时跨参与进程广播，并且梯度在backward() 期间在进程中 allreduced and averaged。

DistributedDataParallel 针对与 NCCL 一起使用进行了优化。

DistributedDataParallel 旨在与上游启动实用程序脚本 torch.distributed.launch 一起使用 --nproc_per_node