Pytorch Multi-GPU原理与实现（单机多卡）

Posted 2023-03-16

参考技术A 最近训练模型时候想要使用使用多GPU运算来提高计算速度，参考一些博客以及自己的动手实验搞懂了Pytorch的Multi-GPU原理。现在稍微整理一下。

通常情况下，多GPU运算分为单机多卡和多机多卡，两者在pytorch上面的实现并不相同，因为多机时，需要多个机器之间的通信协议等设置。

pytorch实现单机多卡十分容易，其基本原理就是：加入我们一次性读入一个batch的数据, 其大小为[16, 10, 5]，我们有四张卡可以使用。那么计算过程遵循以下步骤：

定义模型

定义优化器

实现多GPU

注：模型要求所有的数据和初始网络被放置到GPU0，实际上并不需要，只需要保证数据和初始网路都在你所选择的多个gpu中的第一块上就行。

pytorch multi-gpu train

记录一下pytorch如何进行单机多卡训练：

官网例程：https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html

下面以一个例子讲解一下，例如现在总共有8张卡，在第5、6、7三张卡上进行训练；

step 1:可视化需要用到的GPU

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "5 , 6 , 7"

device = torch.device("cuda:0")  #注意多卡训练的时候，默认都先将model和data先保存在id:0的卡上（即实际的第5块卡），然后model的参数会复制共享到其他卡上，data也会平分成若干个batch到其他卡上（所以第一块卡上稍微耗费一点显存）；

device_ids = [0 , 1 , 2] #注意device_ids必须从0开始，代码中的所有的device id都需要从0开始（这里的0代表第5块卡，1代表第6块卡，类推）；

step 2:利用DataParallel对Model类进行封装

model = nn.DataParallel(model , device_ids = device_ids)

model.to(device)

step 3:

data.to(device)  #id:0卡上的数据再被平分成若干个batch到其他卡上

注意：晚上还有一些例程，需要对optimizer和loss利用DataParellel进行封装，没有试验过，但上面方法是参考官网例程，并经过实操考验；