持续创作，加速成长！这是我参与「掘金日新计划 · 6 月更文挑战」的第4天，点击查看活动详情

单显卡

先看单显卡如何实现训练：第一步：检查环境内是否有cuda环境，如果有则将device 设置为cuda，如果没有则设置为cpu

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

第二步：定义模型，并将模型放入device。

model = mobilenet_v3_large(pretrained=True)
model.to(DEVICE)

第三步：将数据放入device中

 for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)

第四步：保存模型

torch.save(model, 'model_.pth')

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

多显卡

多卡采用DP的方式，DP是DataParallel的缩写。它使用一个进程来计算模型参数，然后在每个批处理期间将分发到每个GPU，然后每个GPU计算各自的梯度，然后汇总到GPU0中进行求平均，然后由GPU0进行反向传播更新参数，然后再把模型的参数由GPU0传播给其他的GPU。

特点：（1）broadcast 的是模型的参数，因此速度慢，效率低（2）操作简单

缺点：GPU利用率通常很低。nn.DataParallel要求所有的GPU都在同一个节点上（不支持分布式），而且不能使用Apex进行混合精度训练。

第一步：检查环境内是否有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中。数据默认加载到第一张显卡上。

 for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)

第四步：保存模型，判断是否是多张显卡训练出来的模型。如果是多张显卡训练的模型则要保存model.module。

 if isinstance(model, torch.nn.DataParallel):
                torch.save(model.module, 'model_.pth')
            else:
                torch.save(model, 'model_.pth')

在这里插入图片描述

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

单显卡

多显卡