转换onnx_PyTorch 到 ONNX 到 Tensorflow 到 TFLite的坑

189 阅读 0 评论 125 点赞

我是靠谱客的博主害怕纸飞机，这篇文章主要介绍转换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 能处理的节点，如下图所示

以上就是害怕纸飞机最近收集整理的关于转换onnx_PyTorch 到 ONNX 到 Tensorflow 到 TFLite的坑的全部内容，更多相关转换onnx_PyTorch内容请搜索靠谱客的其他文章。

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