pytorch指定用多张显卡训练_关于PyTorch多GPU训练
最近在帮导师赶论文的一个Baseline模型，由于效率要求不得已必须把PyTorch多GPU的功能用起来了。但由于官网的tutorial里的描述相当简单，我学东西又比较想当然，于是像一只无头苍蝇一样撞来撞去，抓住能抓住的一切稻草进行改动，还是没有成功。并且由于课业繁忙，事情也搁置了很久，直到昨晚被高老师愤怒批评。没了良乡的助教，今天算是空出了一整个下午的时间，我重新思考，寻找线索，终于把训练部分的代码跑了起来。
我遇到的主要问题
我的目标是将一个能够在单GPU上完美运行的代码升级为多GPU运行的版本。由于作者在写代码的时候并没有考虑到多GPU的使用情景，因此修改起来要更麻烦一点。
PyTorch的DataParallel其实是把数据分到各个GPU上边，进行运算，然后返回到一个GPU上进行汇总，得到最终梯度的过程。新版本的PyTorch加入了ModelParallel的功能，即将模型分配到各个GPU上，加强了并行性，缓解了显存占用不均衡的问题。但我目前还未仔细研究，因此暂且不表。
需要注意，我目前也仅完成了模型训练部分的代码修改。根据我得到的信息，在测试阶段读取模型的时候，代码依然需要进行修改，我会在那时进行新的update。
具体步骤
注：这一节里所有被注释的代码代表单GPU运算时的代码。
Step-1：设定GPU
通过以下语句确定你要使用的GPU。
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"
1
os.environ["CUDA_VISIBLE_DEVICES"]="1,2,3"
Step-2: 用DataParallel模块对模型进行包装
这里一般直接定位代码里类似model.cuda()或model.to(device)等代码，加入以下内容。
Python
if options["cuda"]:
# model.cuda()
model = nn.DataParallel(model, device_ids=[0, 1, 2])
model.to(torch.device("cuda:0"))
1
2
3
4
ifoptions["cuda"]:
# model.cuda()
model=nn.DataParallel(model,device_ids=[0,1,2])
model.to(torch.device("cuda:0"))
我在这里踩到了坑，网上的一些文章，只提到了使用model = nn.DataParallel(model, device\_ids=[0, 1, 2])这句话来包装model，但事实上这句话仅仅是把模型进行了包装，并没有把其中的变量和内容放到GPU上，因此model.to(torch.device(“cuda:0”))是必需的。如果没有加第二句话，模型会报错:
module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device cpu.
1
modulemusthaveitsparametersandbuffersondevicecuda:0(device_ids[0])butfoundoneofthemondevicecpu.
此外要注意，这里的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进行处理，我这里的情况是
Python
consts["batch_size"] = 5 if options["is_debugging"] else TRAINING_DATASET_CLS.BATCH_SIZE * max(1, torch.cuda.device_count())
consts["testing_batch_size"] = TESTING_DATASET_CLS.BATCH_SIZE * max(1, torch.cuda.device_count())
1
2
3
consts["batch_size"]=5ifoptions["is_debugging"]elseTRAINING_DATASET_CLS.BATCH_SIZE*max(1,torch.cuda.device_count())
consts["testing_batch_size"]=TESTING_DATASET_CLS.BATCH_SIZE *max(1,torch.cuda.device_count())
如果这一步没有做，其实也没什么问题，只要你的batch size可以整除N就行，但速度其实没啥变化，因为每次仍然只处理原来数量的batch，并没有变得更多。你会发现gpustat中你的内存实际上是原来的数据平均分在三张卡上，但如果你乘了N，那N张卡上都会有你之前一张卡上的数据量，也就是大概N倍的提速。
Step-4: 对Loss进行修改
由于我们进行了多GPU训练，因此，得到的loss其实是多个GPU返回的loss列表。因此我们需要对其进行平均操作，才能顺利的进行backward梯度计算。具体操作如下：
Python
# cost.backward()
cost.mean().backward()
# cost = cost.item()
cost = cost.mean().item()
1
2
3
4
# cost.backward()
cost.mean().backward()
# cost = cost.item()
cost=cost.mean().item()
这一步如果没做，出现的情况会是模型可以正常进入训练状态，但是会卡在第一步不动，然后显卡占用率100%。我第一次就是没做这一步，然后放程序跑了两天，但并没有任何效果。
Step-5: 对模型代码进行修改
这是比较棘手的一步，我最开始就是在这里搞蒙了，然后全部失败。
首先，我们进行了Step-2的操作之后，我们的模型，以及其继承、在init函数中引用的所有模型，就全部被自动分配到GPU上了，因此，任何一个模型代码的self成员变量，都不能进行.cuda()或.to(device)操作。因为这样会造成重复放到GPU上边，在只有一个GPU的时候不论怎么放都是那个GPU，多个GPU时情况就有了明显不同。 如果不小心对成员变量进行了.cuda()操作，会出现如下错误：
忘记具体表述了。。总之内容是运算符的两个变量没在一个GPU上边，没法计算。
1
忘记具体表述了。。总之内容是运算符的两个变量没在一个GPU上边，没法计算。
然而，如果不是成员变量，而是在forward函数中新定义的临时变量，如为了做padding而定义的0矩阵，那么则需要对其进行.cuda()或.to(device)操作，否则会出现如下错误
RuntimeError: Expected object of backend CUDA but got backend CPU for argument #2
1
RuntimeError:ExpectedobjectofbackendCUDAbutgotbackendCPUforargument#2
或
RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2
1
RuntimeError:ExpectedobjectofbackendCPUbutgotbackendCUDAforargument#2
总结
以上基本就是我在这两个礼拜里边遇到的所有问题，回想起来其实并不难，理解之后我从原始代码重新修改了一遍，也没花太久。但是由于对框架的理解不清、过于想当然以及过滤信息的能力依然较差，浪费了很多时间。这里还有读取模型部分的代码等待处理，等我搞定之后会再进行更新。