【pytorch】pytorch 转 onnx 踩坑记录

最新推荐文章于 2025-03-17 20:54:58 发布

sxxblogs

最新推荐文章于 2025-03-17 20:54:58 发布

阅读量1k

点赞数 1

分类专栏： python 姿态估计文章标签： pytorch 深度学习 python

本文链接：https://blog.youkuaiyun.com/sunnyblogs/article/details/120054349

版权

本文记录了在将PyTorch模型转换为ONNX格式过程中遇到的三个主要问题及解决方案。第一个问题是由于PyTorch版本不支持字典输出，解决方法是将模型输出更改为list或tuple。第二个问题是在符号操作中遇到KeyError，解决方法是手动添加缺失的remainder定义。第三个错误是ONNX推理时的版本不兼容问题，可能是ONNX版本过高或version参数设置不当，目前问题尚未解决。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import io
import torch
import torch.onnx
import numpy as np
from model.detnet import detnet

'''onnx的输入为array'''
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")

def test():
  model = detnet()
  model = model.to(device)
  
  pthfile = './ckp_detnet_137.pth'
  
  check_point  = torch