Pytorch中多GPU训练

参考

https://www.jianshu.com/p/0bdf846dc1a2

https://oldpan.me/archives/pytorch-to-use-multiple-gpus

前言

在数据越来越多的时代，随着模型规模参数的增多，以及数据量的不断提升，使用多GPU去训练是不可避免的事情。Pytorch在0.4.0及以后的版本中已经提供了多GPU训练的方式，本文简单讲解下使用Pytorch多GPU训练的方式以及一些注意的地方。

这里我们谈论的是单主机多GPUs训练，与分布式训练不同，我们采用的主要Pytorch功能函数为DataParallel而不是DistributedParallel，后者为多主机多GPUs的训练方式，但是在实际任务中，两种使用方式也存在一部分交集。

使用方式

使用多卡训练的方式有很多，当然前提是我们的设备中存在两个及以上的GPU：使用命令nvidia-smi查看当前Ubuntu平台的GPU数量(Windows平台类似)，其中每个GPU被编上了序号：[0,1]：

在我们设备中确实存在多卡的条件下，最简单的方法是直接使用torch.nn.DataParallel将你的模型wrap一下即可：

net = torch.nn.DataParallel(model)

这时，默认所有存在的显卡都会被使用。

如果我们机子中有很多显卡(例如我们有八张显卡)，但我们只想使用0、1、2号显卡，那么我们可以：

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

或者这样：

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0,1,2]))
net = torch.nn.DataParallel(model)
# CUDA_VISIBLE_DEVICES 表示当前可以被python环境程序检测到的显卡

很简单的操作，这样我们就可以比较方便地使用多卡进行训练了。