pytorch显卡内存随训练过程而增加_pytorch apex 混合精度训练和horovod分布式训练...

转载请注明出处：https://zhuanlan.zhihu.com/p/98034129

如果你基于pytorch训练模型，然后，你想加快训练速度，增大batch_size,或者，你有一台配置多张显卡的机器，还是说你有多台带显卡机器，你想利用起来，分布式训练你的模型，那这篇文章对你有点用。

基于以上的需求，我趟了一遍，记录下我遇到的坑都有哪些，怎么跨过去。

先看一下我主要的工具：anaconda,apex,horovod。就这三个。

1. Conda 环境

接下来的工作，全部在conda环境内完成，所以安装一个最新版的conda，然后创建一个新的环境，并且激活这个环境。

在这个全新的环境里面，安装模型训练需要的所有package。

三种配置conda新环境的方法：

a.最最简单的办法，直接拷贝 path_to_conda/anaconda3/envs 下面的环境，移植到新机器对应环境下。但是这个方法有前提条件：首先是，最好新旧机器的硬件配置、系统版本是一样的；第二是环境内的package的安装过程不需要配置额外参数。

b.导出conda 环境，在新机器上创建。显然这种方式的前提和第一种相似。

conda env export > environment.yaml  #conda导出已有环境

conda env create -f environment.yaml #在新机器上创建

c.当环境无法移植进来时，只能一个一个安装

安装好环境后建议检查以下cuda和pytorch版本是否匹配，在python脚本内使用以下方法检查:

import torch

torch.cuda.is_available()

2. 混合精度训练和单机多卡分布式训练

Apex 是Nvidia开发维护的工具，使用简便的API实现混合精度训练和分布式训练 根据以下链接安装apex：

https://github.com/NVIDIA/apex#quick-start​github.com <