pytorch 并行训练 pytorch gpu并行

目录

• 写在前面

• 如何使用多GPU并行训练
• 踩坑集锦

• 1.GPU内存不够
• 2.数据集的设置

• 写在后面

写在前面

最近笔者在使用超算的多GPU并行训练Swin Transformer模型。在训练大模型和使用大规模数据集的时候，为提升训练速度，往往都需要多GPU并行训练。因此笔者写下此文，作为曾经踩坑的记录，并帮助需要用到多GPU训练的朋友们解决自己的问题。

如何使用多GPU并行训练

关于单GPU的使用，由于难度不大，笔者不多做赘述，这里介绍多GPU的并行训练方法。
首先，我们将待会需要调用的GPU设置为可见状态。以4块GPU为例，在代码的最顶上加入以下代码：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
ids = [0,1,2,3]

其中第二行代码将四块GPU设置为可见状态，第三行为四块GPU进行编号。之后，我们需要将模型打包，修改如下所示：

model = models.vgg16()
###
###导入网络权重
###
model = torch.nn.DataParallel(model, device_ids=ids).cuda()

首先定义我们需要使用到的模型，这里以一个VGG16模型为例。然后我们导入网络的预训练权重，之后，使用最后一行代码中的DataParallel函数将四个model打包起来，形成一个打包后的model，那么程序在运行的时候，就会在四个GPU上并行运行模型。

这里需要注意了，使用DataParallel函数打包模型的时候，必须已经导入了网络权重，不然就会报错如下图所示。

RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]), but found one of them on device : CPU.

这里的报错意思是，需要在cuda:0上有参数，但是只在GPU上找到参数。乍一看这句话，笔者一开始以为是没有将CPU上面的权重转移到GPU上，但是在之前的代码model = torch.nn.DataParallel(model, device_ids=ids).cuda()里面，.cuda()已经将CPU的参数转移到GPU上了，后来发现，是因为没有先导入权重，再设置并行化代码model = torch.nn.DataParallel(model, device_ids=ids).cuda()，所以，此时根本没有参数！所以再一次强调一下，使用DataParallel函数打包模型的时候，必须已经导入了网络权重。
设置完成以上内容，即可进行多GPU并行训练。但是，可能还是会遇到很多问题，接下来是一些笔者遇到过的问题。

踩坑集锦

1.GPU内存不够

RuntimeError: CUDA out of memory. Tried to allocate ... MiB

这个问题应该是很多人使用GPU的时候都会遇到的问题，很好解决。最简单的方法：改小batch_size的大小即可解决问题。除此之外，还可以在模型验证的时候加上 with torch.no_grad()等方法，还有很多其他方法可以在网上查找。

2.数据集的设置

这是个大坑，也是一个使用多GPU并行训练很可能遇到的问题，笔者训练了七个小时的模型，训练到99%的时候，因为这个坑导致之前的训练前功尽弃。

在多GPU训练的时候，数据集总量$N_{data}$，GPU数量$N_{GPU}$，以及batch_size的数量$N_{batchsize}$，必须满足一定的关系：

$N_{data}/N_{batchsize}= N······N_{left},$

其中$N$是整除的结果，$N_{left}$是余数，$N_{left}$必须等于0。

以笔者的数据集为例，训练集总数$N_{data}$为8607，batch_size的数量$N_{batchsize}$为4，那么$N$就等于2151，余数$N_{left}$等于3。当运行到以下阶段的时候：

报错如下所示：

其中的TypeError: forward() missing 1 required positional argument: 'vclips'不是根源问题，根源在于TypeError: Caught TypeError in replica 3 on device 3. 这个报错什么意思呢？8607个数据集，以4个数据为一组，刚好可以分为2151组数据，余下3个数据。每组的4个数据分配给4个GPU，每个GPU运行一个数据。当运行完前2152组之后，开始分配最后3个数据，可以此时有4个GPU，编号为0，1，2的GPU都分配到了数据，但是编号为3的GPU，没有数据了，所以产生错误！

解决思路：调整数据集的总数、GPU总数和batchsize的大小关系。

写在后面

上面的内容基于使用一台服务器的多GPU进行并行训练，之后遇到有价值的问题，笔者会持续更新。欢迎各位指正和提问。