torch DDP训练-模型保存-加载问题

单GPU保存与加载:

模型保存:
####方法一
state = {'epoch': epoch,
         'model': model.state_dict(),
         'optimizer': optimizer.state_dict(),
         'scheduler': scheduler.state_dict()}
torch.save(state, 'model_path')

####方法二
torch.save(self.model.state_dict(), 'model_path')  # 当前目录

#######################################################################################


模型加载:
####方法一
self.model.load_state_dict(torch.load("model_path")['model'])

####方法二
self.model.load_state_dict(torch.load('model_path'))

DDP模式下模型保存与加载:

模型保存:
                    无module形式
####方法一
state = {'epoch': epoch,
         'model': model.state_dict(),
         'optimizer': optimizer.state_dict(),
         'scheduler': scheduler.state_dict()}
torch.save(state, 'model_path')

####方法二
torch.save(self.model.state_dict(), 'model_path')


                    module形式---建议模式
####方法一
state = {'epoch': epoch,
         'model': model.module.state_dict(),
         'optimizer': optimizer.state_dict(),
         'scheduler': scheduler.state_dict()}
torch.save(state, 'model_path')

####方法二
torch.save(self.model.module.state_dict(), 'model_path')


#######################################################################################


模型加载:  
        未用含"moduel"方式保存, 导致缺失关键“key”:Missing key(s) in state_dict
############### 方法 1: add
model = torch.nn.DataParallel(model)  # 加上module
model.load_state_dict(torch.load("model_path"))

############### 方法 2: remove
model.load_state_dict({k.replace('module.', ''): v for k, v in                 
                       torch.load("model_path").items()})

############### 方法 3: remove
from collections import OrderedDict
state_dict = torch.load("model_path")
new_state_dict = OrderedDict()   # create new OrderedDict that does not contain `module.`
for k, v in state_dict.items():
    name = k.replace('module.', '')
    new_state_dict[name] = v
model.load_state_dict(new_state_dict)

                        含"moduel"方式保存
####方法一
self.model.load_state_dict(torch.load("model_path")['model'])

####方法二
self.model.load_state_dict(torch.load('model_path))

 DDP训练流程:

  • 初始化 

    #初始化使用nccl后端(这个),当然还有别的后端,可以查看官方文档,介绍的比较清晰
    torch.distributed.init_process_group(backend="nccl")
  • 使用DistributedSampler

    DDP并不会自动shard数据如果自己写数据流,得根据torch.distributed.get_rank()去shard数据,获取自己应用的一份如果用Dataset API,则需要在定义Dataloader的时候用DistributedSampler 去shard:

  • 分布式训练

    model=torch.nn.parallel.DistributedDataParallel(model)
  • 参考文章https://zhuanlan.zhihu.com/p/95700549https://zhuanlan.zhihu.com/p/95700549

https://zhuanlan.zhihu.com/p/145427849icon-default.png?t=L892https://zhuanlan.zhihu.com/p/145427849 

DDP训练问题:

1.自定义的模型结构,继承自nn.Module,除__init__(), forward()等重写方法外,在模型结构类内部自定义的一些函数如loss()等,在DDP训练时调用方式需注意

model = YoulModel()
model.loss()  ### 找不到属性loss

#在使用net = torch.nn.DDP(net)之后,原来的net会被封装为新的net的module属性里

model.module.loss() 

  • 14
    点赞
  • 39
    收藏
    觉得还不错? 一键收藏
  • 6
    评论
torch加载训练模型非常简单,只需几行代码即可完成。首先,我们需要导入相应的库和模型。 ```python import torch import torchvision.models as models ``` 接下来,选择一个预训练模型,比如ResNet50。 ```python model = models.resnet50(pretrained=True) ``` 在这个例子中,我们加载了一个在ImageNet数据集上预训练好的ResNet50模型。`pretrained`参数设置为`True`,表示我们要加载训练模型。请确保已经安装了torchvision库,它包含了许多常用的预训练模型。 接下来,可以通过打印模型的结构来查看模型的信息。 ```python print(model) ``` 可以看到模型的结构、参数信息等。 如果我们只需要使用模型进行推断而不需要微调,我们可以将模型设置为评估模式,以节省内存和加快计算速度。 ```python model.eval() ``` 现在我们已经成功加载了预训练模型,可以将输入数据传递给模型进行推断了。 ```python output = model(input) ``` 这里的`input`是输入到模型的数据。输出结果`output`是一个张量,其中包含了模型对输入数据的预测结果。 最后,根据模型的需求进行后处理,比如应用Softmax函数将结果转换为概率分布,或者进行其他操作。 以上就是使用torch加载训练模型的简单步骤。根据需要,我们可以在加载模型后进行微调,修改模型的最后几层,或者使用模型的特征提取层作为新问题的输入。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值