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.

猜你喜欢

转载自blog.csdn.net/znsoft/article/details/128805284