.pt/.pth转onnx踩的坑 小记

收到一个.pt的人脸关键点模型,需要将其转为onnx,很简单的一步,但闹出了乌龙,经过小小折腾解决后将过程记录一下。

torch.onnx.export一行代码即可进行pytorhc模型与onnx的转换,但获得的onnx模型跑视频的时候精度掉的很多,于是打印了一下.pt与.onnx两个的输出进行对比,发现相差较大且毫无规律。

然后就是排查过程,分为以下的心路历程:

1.排查问题的时候,还以为是有一些算子不兼容,看了一下.pt模型的激活函数是hardtanh,而onnx中是clip算子,看了onnx文档后发现是一致的。

下图是hardtanh(从其他博客截的图):

下图是clip算子onnx文档的解释:

 那问题到底出在哪里呢?

 2.以为是torch.onnx版本导致的问题

但切换几个版本后结果还是一样的。

3.模型保存问题 

之前转Onnx的时候,并不需要指定output的格式,只需要给定input的shape即可,而这次export则需要给定example_outputs,所以感觉会不会是.pt模型保存的不对。

于是print一下模型,发现里面的结构全是RecursiveScriptModule和original_name = xxx:

(features): RecursiveScriptModule(
original_name=Sequential
(0): RecursiveScriptModule(
original_name=Conv2dNormActivation
(0): RecursiveScriptModule(original_name=Conv2d)
(1): RecursiveScriptModule(original_name=BatchNorm2d)
(2): RecursiveScriptModule(original_name=ReLU6)

!这不是torchscript的模型吗,怪不得之前加载的时候并不需要提供结构类的定义,我还以为是torch.load(net)加载了参数和结构。 这里又有一个小tips,先偏一下题。。

torch.save的两种方式,torch.save(net,filename)和torch.save(net.state_dict(), filename),前者是整个模型保存,后者是只保存参数,结构需要引用类的定义。这里非常的误导人,因为即使你使用保存全部模型,collections.OrderedDict()是序列化保存了参数和原结构类的路径,也就是说在你torch.load()的时候,如果你的工程或电脑变了,则无法获取模型结构。所以建议大家还是选择第二种保存方式。

然后回到正题,为什么torchscript转onnx会有问题。

因为torchscript模型是pytorch模型经过torch.jit.script或torch.jit.trace得到的,我拿到的这个模型应该是torch.jit.script转换的,所以模型内部的数据类型只有张量,且没有for if while等控制流,不支持python的预处理和动态行为。但好处就是可以连带模型的定义一起保存。

所以要求对方提供一个torch.save(net.state_dict)的模型和模型结构即可成功转成onnx啦。

猜你喜欢

转载自blog.csdn.net/yjcccccc/article/details/131455679