Pytorch的数据并行

在模型训练时，如果有多张GPU，那么一般会使用到数据并行。

数据并行（ data parallelism ）：不同的机器有同一个模型的多个副本，每个机器分配到不同的数据，然后将所有机器的计算结果按照某种方式合并。

在pytorch中，数据并行通常使用torch.nn.DataParallel函数来实现。

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

使用上面的代码后，可以在terminal中使用nvidia-smi来查看显卡信息，会发现第一块卡的现存占用会高一些，这是什么原因导致的呢？查看pytorch官网torch.nn.DataParallel的文档说明，可以看到

CLASS torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

module就是自己定义的模型，device_ids就是GPU的id，output_device表示模型输出的设备位置，一般省略不写，而默认是在device_ids[0]，也就是第一块卡，这样第一块卡的显存会占用比其他卡要更多一些。使用torch.nn.DataParallel时，input数据是并行的，但是output loss不是，每次都会在第一块卡上相加计算，这就造成第一块显卡的负载远大于剩余的显卡。

但使用torch.nn.DataParallel时有个小坑，在官网中有这样一句话

The parallelized module must have its parameters and buffers on device_ids[0] before running this DataParallel module.

就是并行化模块必须在device_ids[0]上具有其参数和缓冲区，在执行torch.nn.DataParallel之前，会首先把其模型的参数放在device_ids[0]上。这样就造成在使用多张GPU时，device_ids必须包含第一张卡，否则会出现模型初始化错误。而解决这个问题，可以通过下面两行代码

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

这样就改变了device的逻辑id。（这两行代码要放在torch.nn.DataParallel之前）

此外，使用torch.nn.DataParallel的模型，在加载模型时也有些特别，通常加载模型使用

model.load_state_dict(checkpoint['state_dict'])

而使用了数据并行的模型，加载模型的命令为

model.module.load_state_dict(checkpoint['state_dict'])

参考资料：
https://zhuanlan.zhihu.com/p/...