1. Toy模板网首页
  2. > Toy博客

将pytorch的pth文件固化为pt文件

7月前 作者：yuanlulu 分类：Toy博客 阅读(34) 违法举报

这篇具有很好参考价值的文章主要介绍了将pytorch的pth文件固化为pt文件。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

说明

我参考了一个开源的人像语义分割项目mobile_phone_human_matting，这个项目提供了预训练模型，我想要将该模型固化，然后转换格式后在嵌入式端使用。

该项目保存模型的代码如下：

lastest_out_path = "{}/ckpt_lastest.pth".format(self.save_dir_model)
        torch.save({
            'epoch': epoch,
            'state_dict': model.state_dict(),
        }, lastest_out_path)

转换代码

上面代码保存了state_dict, 所以保存的文件中是不含模型结构的，固化时需要从代码构造网络结构。好在项目是完全开源，将原项目下的model目录拷贝过来就行。
另外不能忘记调用eval() 来固化参数。

完整的转换代码如下：

import torch
from model import segnet

ckptfile="./ckpt_lastest.pth"
savedfile="./human_seg.pt"

model = segnet.SegMattingNet()    

device = torch.device('cpu')
ckpt = torch.load(ckptfile, map_location=device )
model.load_state_dict(ckpt['state_dict'])

model.eval() #这一步会将参数固化，不能省。否则会报AssertionError('batchnorm with training is not support. Please set model.eval() before export.')

x = torch.rand(1,3,256,256)
ts = torch.jit.trace(model, x)
ts.save(savedfile)

参考资料

mobile_phone_human_matting

pytorch训练的.pth模型格式转换文章来源地址https://www.toymoban.com/news/detail-687127.html

到了这里，关于将pytorch的pth文件固化为pt文件的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处： 如若内容造成侵权/违法违规/事实不符，请点击违法举报进行投诉反馈，一经查实，立即删除！

分享到：
领支付宝红包 赞助服务器费用

相关文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包，优惠每天领

二维码1

领取红包

二维码2

领红包