前言

在数据越来越多的时代，随着模型规模参数的增多，以及数据量的不断提升，使用多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环境程序检测到的显卡

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

另一种方法

在前言中提到过，另一种方法DistributedParallel，虽然主要的目标为分布式训练，但也是可以实现单主机多GPU方式训练的，只不过比上一种方法稍微麻烦一点，但是训练速度和效果比上一种更好。

为什么呢？

请看官方相关介绍：

nccl backend is currently the fastest and highly recommended backend to be used with Multi-Process Single-GPU distributed training and this applies to both single-node and multi-node distributed training

好了，来说说具体的使用方法(下面展示一个node也就是一个主机的情况)为：

python -m torch.distributed.launch --nproc_per_node=你的GPU数量
               YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
               arguments of your training script)

上述的命令和我们平常的命令稍有区别，这里我们用到了torch.distributed.launch这个module，我们选择运行的方式变换为python -m，上面相当于使用torch.distributed.launch.py去运行我们的YOUR_TRAINING_SCRIPT.py，其中torch.distributed.launch会向我们的运行程序传递一些变量。

为此，我们的YOUR_TRAINING_SCRIPT.py也就是我们的训练代码中这样写(省略多余代码，只保留核心代码)：

import torch.distributed as dist
# 这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
parser.add_argument("--local_rank", type=int, default=0) 

def synchronize():
    """
    Helper function to synchronize (barrier) among all processes when
    using distributed training
    """
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return
    world_size = dist.get_world_size()
    if world_size == 1:
        return
    dist.barrier()


## WORLD_SIZE 由torch.distributed.launch.py产生 具体数值为 nproc_per_node*node(主机数，这里为1)
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

is_distributed = num_gpus > 1

if is_distributed:
    torch.cuda.set_device(args.local_rank)  # 这里设定每一个进程使用的GPU是一定的
    torch.distributed.init_process_group(
        backend="nccl", init_method="env://"
    )
    synchronize()

# 将模型移至到DistributedDataParallel中，此时就可以进行训练了
if is_distributed:
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[args.local_rank], output_device=args.local_rank,
        # this should be removed if we update BatchNorm stats
        broadcast_buffers=True,
    )

# 注意，在测试的时候需要执行 model = model.module 

我们要注意，上述代码在运行的过程中产生了很多个，具体多少个取决你GPU的数量，这也是为什么上面需要torch.cuda.set_device(args.local_rank)设定默认的GPU，因为torch.distributed.launch为我们触发了n个YOUR_TRAINING_SCRIPT.py进程，n就是我们将要使用的GPU数量。

有一点想问的，我们每次必须要使用命令行的方式去运行吗？当然也可以一键解决，如果我们使用Pycharm，只需要配置成下面这样就可以了：

单显卡与DataParallel多显卡训练对比

最近两天训练一个魔改的mobilenetv2+yolov3，同样的优化方法同样的学习率衰减率，所有的参数都相同的情况下，发现单显卡训练的方式竟然比多显卡训练的方式收敛更快。

配置为两张1080Ti，使用Pytorch的版本为1.0.0。下图红线为使用一张1080Ti训练的情况，蓝线为使用两张1080Ti训练的情况，batchsize每张显卡设置为10，也就是说，使用两张显卡训练时的batchsize为单张显卡的两倍，同一个step时，双卡走的步数为单卡步数的两倍。这里使用的多卡训练方式为DataParallel。

但是下图可以看到，在双卡相同step的情况下，虽然红色曲线的损失相较蓝色下降的稍微慢一些，但是到了一定时候，两者的损失值会相交(此时未达到最低损失点),也就是说使用双卡和单卡训练时候loss损失收敛的速度是一样的。

更奇怪的是，下图中，在验证集中，单显卡虽然没有双显卡的准确度曲线增长迅速，但是到了某一点，单显卡的曲线会超过双显卡训练的精度，也就是说，单卡训练在前期没有双卡训练效果显著，但是到了训练中期效果就会优于双卡。

(上述两个图为训练早期和中期的展示，并没有完全训练完毕)关于为什么会这样的情况，有可能是因为训练中期所有的激活值更新幅度不是很明显(一般来说，权重值和激活值更新幅度在训练前期比较大)，在不同GPU转化之间会损失一部分精度?。当然这仅仅是猜测，博主还没有仔细研究这个问题，待有结论时会在这里进行更新。

注意点

多GPU固然可以提升我们训练的速度，但弊端还有有一些的，有几个我们需要注意的点：

• 多个GPU的数量尽量为偶数，奇数的GPU有可能会出现中断的情况
• 选取与GPU数量相适配的数据集，多显卡对于比较小的数据集来说反而不如单个显卡训练的效果好
• 多GPU训练的时候注意机器的内存是否足够(一般为使用显卡显存x2)，如果不够，建议关闭pin_memory(锁页内存)选项。
• 采用DistributedDataParallel多GPUs训练的方式比DataParallel更快一些，如果你的Pytorch编译时有nccl的支持，那么最好使用DistributedDataParallel方式。

关于什么是锁页内存：

pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

相关资料：
https://blog.csdn.net/tfcy694/article/details/83270701
https://discuss.pytorch.org/t/what-is-the-disadvantage-of-using-pin-memory/1702

总结一句就是，如果机子的内存比较大，建议开启pin_memory=Ture，如果开启后发现有卡顿现象或者内存占用过高，此时建议关闭。

参考文章：

https://pytorch.org/docs/stable/notes/cuda.html#cuda-nn-dataparallel-instead
https://github.com/facebookresearch/maskrcnn-benchmark