pytorch多卡训练模型

DP方法

基本原理

也即是：数据拆分，模型不拆分的思路。

输入数据自动分割成若干小批量，并将模型复制到指定的多个 GPU 上。所有 GPU 计算得到的梯度会在主设备（device_ids 列表中的第一个 GPU）进行汇总后更新模型参数。这时的batch size应该是单块显卡上的batch size乘以显卡数量。

适用场景

适用于模型较小的场景。 但会造成主设备负载过高。

使用方法

#------------------------- 初始化device变量------------------------- 
device = torch.device('cuda:1' if torch.cuda.is_available() else "cpu")

#-------------- 修改batch size为所有显卡训练的batch size总和 --------------
train_dataloader = DataLoader(train_dataset, 8, True)

# ------------------------- 模型初始化、数据并行 -------------------------
model = model(num_class=3)

#调用DataParallel类，传参数。第二个参数为一个list，为参与并行计算的显卡的id，从0开始
#要注意list的第一个数字要和初始化device的显卡号对应一致，否则会报错。
model = model.to(device)  # 先移到主GPU
model = torch.nn.DataParallel(module=model, device_ids=[1, 0, 2, 3])  # 然后包装

主要是上面三个部分要做改动。整体比较简单。

错误记录

1. 调用DataParallel类，传参数。

unet = torch.nn.DataParallel(module=unet, device_ids=[1, 0, 2, 3])。

第二个参数为一个list，为参与并行计算的显卡的id，id从0开始。要注意list的第一个数字要和

device = torch.device('cuda:1' if torch.cuda.is_available() else "cpu")

的显卡号对应一致，否则会报错。

2. train_dataloader = DataLoader(train_dataset, batch size = 8, True)

这里的batch size 应该是所有参加训练的显卡的batch size的总和。因为DataParallel 会将输入 batch 按照设备数量进行切分。如果太小会导致并不是所有的指定的显卡都会参与训练。

DDP方法

目前实验室配置以满足需求，暂不考虑DDP方法。DP方法已经能够完成日常学习。后期若要训练大模型会再学习该方法。