pytorch指定用多张显卡训练_Pytorch中多GPU训练指北

前言

在数据越来越多的时代，随着模型规模参数的增多，以及数据量的不断提升，使用多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])

或者这样：