pytorch 多GPU训练总结（DataParallel的使用）

参考：主页 - PyTorch中文文档

前言：

博主最近搭建网络的时候，需要调用不同的GPU，实现训练的加速，有时间会出现显卡现存分布不均的情况，有时间有的显卡温度特别高。

博客持续更新：

一更：2022.09.01 DP模式见本文，使用最少的代码实现

二更：2022.10.26，DDP可参看，修改的地方稍多，但是效率最高。Pytorch DDP原理及配置[最大限度提高GPU利用率]_两只蜡笔的小新的博客-CSDN博客

Multi-GPU layers

class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)[source]

在模块级别上实现数据并行。

此容器通过将mini-batch划分到不同的设备上来实现给定module的并行。在forward过程中，module会在每个设备上都复制一遍，每个副本都会处理部分输入。在backward过程中，副本上的梯度会累加到原始module上。

batch的大小应该大于所使用的GPU的数量。还应当是GPU个数的整数倍，这样划分出来的每一块都会有相同的样本数量。

请看: Use nn.DataParallel instead of multiprocessing

除了Tensor，任何位置参数和关键字参数都可以传到DataParallel中。所有的变量会通过指定的dim来划分（默认值为0）。原始类型将会被广播，但是所有的其它类型都会被浅复制。所以如果在模型的forward过程中写入的话，将会被损坏。

参数说明：

• module – 要被并行的module
• device_ids – CUDA设备，默认为所有设备。
• output_device – 输出设备（默认为device_ids[0]）

例子：

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