torch.jit保存，加载模型

最新推荐文章于 2024-08-05 12:35:52 发布

AI算法网奇

最新推荐文章于 2024-08-05 12:35:52 发布

阅读量7.7k

点赞数

分类专栏： pytorch知识宝典文章标签：计算机视觉 opencv 前端

本文链接：https://blog.csdn.net/jacke121/article/details/115539086

版权

pytorch知识宝典专栏收录该内容

499 篇文章 226 订阅 ¥29.90 ¥99.00

订阅专栏

超级会员免费看

在尝试使用torch.jit加载模型时遇到Double与Long类型不匹配的错误，这可能由于JIT不支持某些条件判断。YOLOv5的导出过程可能不同。另外，运行代码时出现CUDA库未加载的RuntimeError，可能原因是缺少-Wl,--no-as-needed链接参数导致CUDA库未被正确加载。解决方案包括检查和调整链接参数确保CUDA库被正确引用。" 108577131,9792230,西门子PLC1200 TCP通讯测试与调试,"['PLC编程', '工业网络', '自动化技术', '西门子PLC', 'TCP/IP通信']

摘要由CSDN通过智能技术生成

ori_backbone.eval()
x = torch.randn(1,3,112,112).cuda()
traced_cell = torch.jit.trace(ori_backbone, (x))
torch.jit.save(traced_cell, "jit_checkpoint.pth")


loc=torch.device('cpu')
loc=torch.device('cuda:0')

model=torch.jit.load(spath[0], map_location=loc)
model.eval()

导出jit报错：

Double did not match Long

原因：可能是jit不支持条件判断。

解决方法：

self.model.model[-1].export=True

self.model.eval()
self.model=self.model.float()

self.model.model[-1].export=True
x = torch.randn(1, 3, 544, 960).float().cuda()
traced_cell = torch.jit.trace(self.model, (x