NVIDIA训练深度学习模型加速:APEX库

最近在跑目标检测和图像分类模型,发现很多时候教程里需要 安装apex库,于是我就去网上搜索一下这个,发现apex大有来头;

官方:

https://nvidia.github.io/apex/amp.html

https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html

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

使用理由

使用精度低于32位浮点的数值格式有许多好处。首先,它们需要更少的内存,从而能够训练和部署更大的神经网络。其次,它们需要较少的内存带宽,从而加快数据传输操作。第三,数学运算在降低精度方面运行得更快,特别是在具有TensorCore支持的GPU上。混合精度训练(Mixed Precision Training)实现了所有这些好处,同时确保与完全精度训练相比,不会丢失特定任务的准确性。它这样做的方法是识别需要完全精度的步骤,只对这些步骤使用32位浮点,而在其他地方使用16位浮点。

在PyTorch中的使用:
首先需要安装其apex库(我还没装过),其github地址:https://github.com/NVIDIA/apex。
然后在训练的脚本(代码)中简单添加几句就可以了

from apex import amp

amp.init()
amp.init_trainer(trainer)
with amp.scale_loss(loss, trainer) as scaled_loss:
   autograd.backward(scaled_loss) 

APEX的配置

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

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" 

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

pip install -v --disable-pip-version-check --no-cache-dir ./

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

查看能否正确导入apex:

from apex import amp

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注