Pytorch模型转换为onnx或ncnn的方法兼谈pytorch模型编写规范

最近在将一个其它人编写的pytorch模型转换到ncnn 格式, 转换路径为:

pytorch-> torchscript->pnnx, 中间有副产品onnx模型。

使用torch.jit.trace_module( model, {"entrypoint",[parameters]}) 来记录模型,entrypoint在大多数模型中是 forward, 如果有定制,可以修改。

使用torch.jit.save来保存模型。

影响成功的主要问题是:在模型的forward函数中不能使用条件语句,否则 会影响trace。此时的办法就是去掉里面的条件语句。

编写规范:

1. forward中不要使用条件语句

2. 如果要用script函数编译转换模型,不要使用**args这样的方式来定义函数参数,要写全。

对于trace或script保存的模型,可以用pnnx来转换成ncnn模型,同时它会生成onnx模型,解决onnx转换问题。当然你也可以直接使用torch.onnx.export来转换模型到onnx.

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值