已知模型结构的 ONNX 模型转 PyTorch 模型

缘由

我们已经有很多 ONNX 模型转 PyTorch 模型的工具了, 但是根据我的个人体验, 模型稍微复杂一点, 这些工具就寄了.

在我的一个任务中, 我有一个已知结构, 实现完全相同, 但是比较复杂的 ONNX 模型, 出于一些原因, 我需要把它转成对应这个实现的 Torch 模型, 这时现有的工具并不能很好的完成任务–这些模型都不能利用现有的模型实现, 而且应付不了复杂的模型.

过程

想法

由于我们的目的是逆向一个结构已知, 而且已经获得其实现的模型, 那我们自然可以用一些通用工具用不了的思路. 比如说, 我可以大胆猜测 ONNX 模型生成时使用了一种比较良好的方式, 使得这个转换过程是稳定, 一致的.

观察

经过对要转换的 ONNX 模型和 Torch 模型(随机初始化出来的) 的观察, 我发现:

• 两个模型的权重数量相近, 具体来说, 是 Encoder 部分差2, 而 Decoder 部分相等, 而 Encoder 部分差的那两个是和池化相关的, 应该是原模型没有开启.

• ONNX 中的权重大多数都有可读的名字, 但是少数的名字变成了如 onnx::MatMul_1234 这种不可读的形式.

这里比较麻烦的就是这些名字丢失的权重, 不过因为权重的数量是相近的, 我还是猜测这些权重之间是一一对应的.

另外由于部分权重的名字仍然可见, 我关于 ONNX 模型生成方式的猜测得到了印证, 在我的例子中, 具体来说, 就是

optimum-cli export onnx --task image-to-text --model torch_model onnx_model

这个工具还是比较友好的, 他转出来的 ONNX 都是一致的, 这就方便了我们更直接地得知导出过程中到底发生了什么.

实验

我首先设法构造了这样的一个模型, 它的权重是以 0.1 递增的数字, 然后将其导出为 ONNX.

这时我发现 ONNX 模型中名称不可读的权重果然都与 PyTorch 中的某些权重对应, 前面提到的差 2 的权重无关紧要.

但是有一个问题是, 这些名字丢失的权重, 虽然内部的值(全都是0.x)对得上 PyTorch 模型中的某个权重, 但是有一部分的 shape 却对不上, 具体来说, 是一些矩阵的大小转置了.

为了确认大小转置了是不是里面的值也确实转置了, 我又构造了另一个模型, 这个模型中的权重都是不对称的, 第一行全都是某个值, 其余行全都是另一个值, 这两个值递增且不相等.

这次的结果确认了我的想法, 那些名字丢失的矩阵就是被转置了. 关于为什么要转置, 我想可能是因为某种计算过程的改变:

$A B = (B^T A^T)^T$

总之我们找到了 ONNX 和 PyTorch 模型的权重之间的对应关系.

结果

结果证明了我的想法是对的. 根据这个原理我编写了一个转换程序, 转换得到的模型与原模型的行为完全一致.

转换程序已开源, 我并没有把这个写成一个非常通用的工具, 只是以本文的这个例子走一遍流程, 如果真的有同样的需求应该也能基本复制.

留言评论