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

Posted AI浩

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了pytorch 多卡训练,模型保存的一些问题相关的知识,希望对你有一定的参考价值。

单显卡

先看单显卡如何实现训练:
第一步:检查环境内是否有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')

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

多显卡

第一步:检查环境内是否有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 多卡训练,模型保存的一些问题的主要内容,如果未能解决你的问题,请参考以下文章

Pytorch使用分布式训练,单机多卡

Pytorch使用分布式训练,单机多卡

pytorch单机多卡训练

Pytorch Multi-GPU原理与实现(单机多卡)

每天讲解一点PyTorch 18多卡训练torch.nn.DataParallel

每天讲解一点PyTorch 18多卡训练torch.nn.DataParallel