关于PyTorch多GPU训练

最近在帮导师赶论文的一个Baseline模型,由于效率要求不得已必须把PyTorch多GPU的功能用起来了。但由于官网的tutorial里的描述相当简单,我学东西又比较想当然,于是像一只无头苍蝇一样撞来撞去,抓住能抓住的一切稻草进行改动,还是没有成功。并且由于课业繁忙,事情也搁置了很久,直到昨晚被高老师愤怒批评。没了良乡的助教,今天算是空出了一整个下午的时间,我重新思考,寻找线索,终于把训练部分的代码跑了起来。

我遇到的主要问题

我的目标是将一个能够在单GPU上完美运行的代码升级为多GPU运行的版本。由于作者在写代码的时候并没有考虑到多GPU的使用情景,因此修改起来要更麻烦一点。
PyTorch的DataParallel其实是把数据分到各个GPU上边,进行运算,然后返回到一个GPU上进行汇总,得到最终梯度的过程。新版本的PyTorch加入了ModelParallel的功能,即将模型分配到各个GPU上,加强了并行性,缓解了显存占用不均衡的问题。但我目前还未仔细研究,因此暂且不表。
需要注意,我目前也仅完成了模型训练部分的代码修改。根据我得到的信息,在测试阶段读取模型的时候,代码依然需要进行修改,我会在那时进行新的update。

具体步骤

注:这一节里所有被注释的代码代表单GPU运算时的代码。

 Step-1:设定GPU

通过以下语句确定你要使用的GPU。

Step-2: 用DataParallel模块对模型进行包装

这里一般直接定位代码里类似model.cuda()或model.to(device)等代码,加入以下内容。
我在这里踩到了坑,网上的一些文章,只提到了使用model = nn.DataParallel(model, device\_ids=[0, 1, 2])这句话来包装model,但事实上这句话仅仅是把模型进行了包装,并没有把其中的变量和内容放到GPU上,因此model.to(torch.device(“cuda:0”))是必需的。如果没有加第二句话,模型会报错:
此外要注意,这里的device_ids是从0开始的最开始选中显卡的相对位置编号。如,我最开始指定了2,4,6,8号显卡,这里将用0,1,2,3来对其进行表示。

Step-3: 更改模型的Batch_Size

由于PyTorch默认每次会把Batch_Size大小的模型放到多个GPU上去,因此每一步运算的数据量其实变成了Batch_Size * N,其中N为分配的GPU个数。所以我们要对batch_size进行处理,我这里的情况是
如果这一步没有做,其实也没什么问题,只要你的batch size可以整除N就行,但速度其实没啥变化,因为每次仍然只处理原来数量的batch,并没有变得更多。你会发现gpustat中你的内存实际上是原来的数据平均分在三张卡上,但如果你乘了N,那N张卡上都会有你之前一张卡上的数据量,也就是大概N倍的提速。

Step-4: 对Loss进行修改

由于我们进行了多GPU训练,因此,得到的loss其实是多个GPU返回的loss列表。因此我们需要对其进行平均操作,才能顺利的进行backward梯度计算。具体操作如下:
这一步如果没做,出现的情况会是模型可以正常进入训练状态,但是会卡在第一步不动,然后显卡占用率100%。我第一次就是没做这一步,然后放程序跑了两天,但并没有任何效果。

Step-5: 对模型代码进行修改

这是比较棘手的一步,我最开始就是在这里搞蒙了,然后全部失败。
首先,我们进行了Step-2的操作之后,我们的模型,以及其继承、在init函数中引用的所有模型,就全部被自动分配到GPU上了,因此,任何一个模型代码的self成员变量,都不能进行.cuda()或.to(device)操作。因为这样会造成重复放到GPU上边,在只有一个GPU的时候不论怎么放都是那个GPU,多个GPU时情况就有了明显不同。 如果不小心对成员变量进行了.cuda()操作,会出现如下错误:

然而,如果不是成员变量,而是在forward函数中新定义的临时变量,如为了做padding而定义的0矩阵,那么则需要对其进行.cuda()或.to(device)操作,否则会出现如下错误

总结

以上基本就是我在这两个礼拜里边遇到的所有问题,回想起来其实并不难,理解之后我从原始代码重新修改了一遍,也没花太久。但是由于对框架的理解不清、过于想当然以及过滤信息的能力依然较差,浪费了很多时间。这里还有读取模型部分的代码等待处理,等我搞定之后会再进行更新。

发表评论

电子邮件地址不会被公开。 必填项已用*标注