参考
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环境程序检测到的显卡
很简单的操作,这样我们就可以比较方便地使用多卡进行训练了。