pytorch 多卡训练，模型保存的一些问题

torch.save(model, ‘model_.pth’)

接下来，我们看看多显卡训练过程和单显卡有何不同

多显卡

==============================================================

第一步：检查环境内是否有cuda环境，如果有则将device 设置为cuda:0，如果没有则设置为cpu。

device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)

第二步：定义模型，并将模型放入device。如果环境中有多张显卡，则使用torch.nn.DataParallel方法加载模型，实现多卡训练。

model = mobilenet_v3_large(pretrained=True)

model.to(DEVICE)

if torch.cuda.device_count() > 1:

print(“Let’s use”, torch.cuda.device_count(), “GPUs!”)

model = torch.nn.DataParallel(model)

第三步：将数据放入device中。数据默认加载到第一张显卡上。</