pytorch模型格式转换

将生成的ckpt_e_50.pth文件转为适合在pc端做推断的.pt文件:

model = UNet(3, 1)
modelname = 'ckpt_e_50.pth'
ckpt = torch.load(opt.pretrain + modelname)
model.load_state_dict(ckpt['state_dict'], strict=False)
model.eval()
example = torch.rand(1, 3, 240, 320)
example1 = torch.rand(1, 1, 30, 40)
traced_script_module = torch.jit.trace(model, (example, example1))
traced_script_module.save("model.pt")

注意,开始保存出的model.pt在pc端C++载入的时候报错,发现torch版本是1.0.0,更新为1.1.0之后转出来的pt文件即可正确读取。

(0)

相关推荐

  • Pytorch入门:Mask R

    Contents 1. 处理数据集 2. Mask R-CNN微调模型 3. 模型的训练及验证 4. 遗留问题(解决后删掉)      通过微调预训练模型Mask R-CNN来完成目标检测及语义分割, ...

  • Open3D面向机器学习的扩展库

    Open3D-ML是Open3D的一个扩展,用于3D机器学习任务.它建立在Open3D核心库之上,并通过机器学习工具对其进行扩展,以进行3D数据处理.此repo集中于语义点云分割等应用程序,并提供可应 ...

  • Yolov5 系列2

    上一篇<Yolov5 系列1- Yolo发展史以及Yolov5模型详解>讲了Yolo的发展历史,这一篇的目的是讲述如何使用Yolo v5训练自定的数据集,并会分析一些常见的选项以及背后的故 ...

  • 深度学习在图像分类中的应用ーー利用 Pytorch 从零开始创建 CNN

    重磅干货,第一时间送达 推荐阅读 31个Python实战项目教你掌握图像处理,PDF开放下载 opencv_contrib扩展模块中文教程pdf,限时领取 引言 本文将解释一个卷积神经网络(CNN)的 ...

  • PyTorch重大更新:将支持自动混合精度训练!

    AI编辑:我是小将 混合精度训练(mixed precision training)可以让模型训练在尽量不降低性能的情形下提升训练速度,而且也可以降低显卡使用内存.目前主流的深度学习框架都开始支持混合 ...

  • 模型部署翻车记:pytorch转onnx踩坑实录

    作者丨nihate 审稿丨邓富城 编辑丨极市平台 极市导读 本文记录了作者在深度学习模型部署是,从pytorch转换onnx的过程中的踩坑记录. >>加入极市CV技术交流群,走在计算机视觉 ...

  • 计算机视觉中自注意力构建块的PyTorch实现

    作者:AI Summer 编译:ronghuaiyang(AI公园) 导读 一个非常好用的git仓库,封装了非常全面的计算机视觉中的自注意力构建块,直接调用,无需重复造轮子了. git仓库地址:htt ...

  • 实用教程详解:用OpenCV的DNN模块部署YOLOv5目标检测

    作者丨nihate 审稿|邓富城 编辑丨极市平台 极市导读 本文中介绍的整套程序只依赖OpenCV库就能正常运行,彻底摆脱了对深度学习框架的依赖.文章讲述了作者在自己编写用OpenCV的dnn模块做Y ...

  • Pytorch量化入门之超分量化(二)

    来源:AIWalker 最近Happy在尝试进行图像超分的INT8量化,发现:pytorch量化里面的坑真多,远不如TensorFlow的量化好用.不过花了点时间终于还是用pytorch把图像超分模型 ...

  • 零基础入门语义分割-Task5 模型训练与验证

    一个成熟合格的深度学习训练流程至少具备以下功能: 在训练集上进行训练,并在验证集上进行验证: 模型可以保存最优的权重,并读取权重: 记录下训练集和验证集的精度,便于调参. 5 模型训练与验证 为此本章 ...

  • 万字长文,60分钟闪电战

    大家好,我是 Jack. 本文是翻译自官方版教程:DEEP LEARNING WITH PYTORCH: A 60 MINUTE BLITZ,一份 60 分钟带你快速入门 PyTorch 的官方教程. ...