赞
踩
收到一个.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啦。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。