pytorch 读取和保存模型参数

只保存参数信息

加载

checkpoint = torch.load(opt.resume)
model.load_state_dict(checkpoint)

保存

torch.save(self.state_dict(),file_path)

这而只保存了参数信息,读取时也只有参数信息,模型结构需要手动编写

保存整个模型

保存
torch.save(the_model, PATH)

加载:
the_model = torch.load(PATH)


有时候会看到加载时

model.load_state_dict(checkpoint['state_dic'])

这是因为checkpoint是一个字典,保存的key可以自己定义。
可以保存除参数信息之外的其它信息,如epoch等。

保存

torch.save({
        'epoch': epoch + 1,
        'arch': args.arch,
        'state_dict': model.state_dict(),
        'best_prec1': best_prec1,
    }, 'checkpoint.tar' )

加载

复制代码
if args.resume:
    if os.path.isfile(args.resume):
        print("=> loading checkpoint '{}'".format(args.resume))
        checkpoint = torch.load(args.resume)
        args.start_epoch = checkpoint['epoch']
        best_prec1 = checkpoint['best_prec1']
        model.load_state_dict(checkpoint['state_dict'])
        print("=> loaded checkpoint '{}' (epoch {})"
              .format(args.evaluate, checkpoint['epoch']))
复制代码

state_dict参考链接:

https://www.cnblogs.com/tingtin/p/13544489.html

 

 

posted on   cltt  阅读(2844)  评论(0编辑  收藏  举报

编辑推荐:
· SQL Server 2025 AI相关能力初探
· Linux系列:如何用 C#调用 C方法造成内存泄露
· AI与.NET技术实操系列(二):开始使用ML.NET
· 记一次.NET内存居高不下排查解决与启示
· 探究高空视频全景AR技术的实现原理
阅读排行:
· 阿里最新开源QwQ-32B,效果媲美deepseek-r1满血版,部署成本又又又降低了!
· 单线程的Redis速度为什么快?
· SQL Server 2025 AI相关能力初探
· AI编程工具终极对决:字节Trae VS Cursor,谁才是开发者新宠?
· 展开说说关于C#中ORM框架的用法!
历史上的今天:
2019-08-22 输出细节
2018-08-22 基本大数问题
< 2025年3月 >
23 24 25 26 27 28 1
2 3 4 5 6 7 8
9 10 11 12 13 14 15
16 17 18 19 20 21 22
23 24 25 26 27 28 29
30 31 1 2 3 4 5

导航

统计

点击右上角即可分享
微信分享提示