转换onnx_PyTorch 到 ONNX 到 Tensorflow 到 TFLite的坑

65 阅读 0 评论 43 点赞

我是靠谱客的博主害怕纸飞机，最近开发中收集的这篇文章主要介绍转换onnx_PyTorch 到 ONNX 到 Tensorflow 到 TFLite的坑，觉得挺不错的，现在分享给大家，希望可以做个参考。

.view()

在 Pytorch 中，经常会对某个具体的 Tensor 做变换的处理。但在 PyTorch 转换到 ONNX 的过程中，例如

fea = fea.view(1, -1, 1, 1)

ONNX 会错误将 view 的纬度扩张变成 Reshape 节点，如下图所示

但是在 Reshape 节点转换到 Tensorflow 的时候，Reshape 节点就变成了 Where 节点和若干节点的组合以防止出现纬度扩张的时候的 Bug，如下图所示

虽然这样子处理貌似稳妥，但是在后续 TF 模型转换到 TFLite 的时候，TFLite 是不支持 Where 节点的。在翻阅 PyTorch Discuss 和 Stackoverflow 的众多讨论后发现大部分人也都是在这里弃坑。

为了避免这个自动转换的问题，这里正确的处理方式为

fea = fea.unsqueeze(-1).unsqueeze(-1)

在修改代码后，我们得到的节点再也不是 Reshape 而是两个 Sequeeze，如下图所示

对应的节点在转换到 Tensorflow的时候，也自动被正确转换成 TFLite 能处理的节点，如下图所示

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。

本图文内容来源于网友提供，作为学习参考使用，或来自网络收集整理，版权属于原作者所有。