pytorch - Pytorch:将 2D-CNN 模型转换为 tflite
问题描述
我想将模型(例如 Mobilenet V2)从 pytorch 转换为 tflite,以便在移动设备上运行它。
有没有人设法做到这一点?
我发现的只是一种使用 ONNX 将模型转换为中间状态的方法。但是,这似乎无法正常工作,因为 Tensorflow 需要 NHWC 通道顺序,而 onnx 和 pytorch 使用 NCHW 通道顺序。
在 github 上有一个讨论,但是在我的情况下,转换工作没有抱怨,直到“冻结的 tensorflow 图模型”,在尝试将模型进一步转换为 tflite 之后,它抱怨通道顺序错误......
到目前为止,这是我的代码:
import torch
import torch.onnx
import onnx
from onnx_tf.backend import prepare
# Create random input
input_data = torch.randn(1,3,224,224)
# Create network
model = torch.hub.load('pytorch/vision:v0.6.0', 'mobilenet_v2', pretrained=True)
model.eval()
# Forward Pass
output = model(input_data)
# Export model to onnx
filename_onnx = "mobilenet_v2.onnx"
filename_tf = "mobilenet_v2.pb"
torch.onnx.export(model, input_data, filename_onnx)
# Export model to tensorflow
onnx_model = onnx.load(filename_onnx)
tf_rep = prepare(onnx_model)
tf_rep.export_graph(filename_tf)
直到这里,所有工作都没有错误(忽略许多 tf 警告)。然后我使用netron(“input.1”和“473”)查找输入和输出张量的名称。
最后,我将我常用的 tf-graph 从 bash 应用到 tf-lite 转换脚本:
tflite_convert \
--output_file=mobilenet_v2.tflite \
--graph_def_file=mobilenet_v2.pb \
--input_arrays=input.1 \
--output_arrays=473
我的配置:
torch 1.6.0.dev20200508 (needs pytorch-nightly to work with mobilenet V2 from torch.hub)
tensorflow-gpu 1.14.0
onnx 1.6.0
onnx-tf 1.5.0
这是我从中得到的确切错误消息tflite
:
Unexpected value for attribute 'data_format'. Expected 'NHWC'
Fatal Python error: Aborted
更新:
更新我的配置:
torch 1.6.0.dev20200508
tensorflow-gpu 2.2.0
onnx 1.7.0
onnx-tf 1.5.0
使用
tflite_convert \
--output_file=mobilenet_v2.tflite \
--graph_def_file=mobilenet_v2.pb \
--input_arrays=input.1 \
--output_arrays=473 \
--enable_v1_converter # <-- needed for conversion of frozen graphs
导致另一个错误:
Exception: <unknown>:0: error: loc("convolution"): 'tf.Conv2D' op is neither a custom op nor a flex op
更新:
这是通过 netron 加载的 mobilenet v2 的 onnx 模型:
这 是我转换后的 onnx 和 pb 文件的 gdrive 链接
解决方案
@Ahwar使用 Google Colab notebook发布了一个很好的解决方案。
它用
torch 1.5.0+cu101
torchsummary 1.5.1
torchtext 0.3.1
torchvision 0.6.0+cu101
tensorflow 1.15.2
tensorflow-addons 0.8.3
tensorflow-estimator 1.15.1
onnx 1.7.0
onnx-tf 1.5.0
转换工作正常,模型可以在我的电脑上测试。然而,当将模型推送到手机时,它仅在 CPU 模式下工作,并且比直接在 tensorflow 中创建的相应模型慢得多(几乎 10 倍)。GPU模式在我的手机上不起作用(对比直接在tensorflow中创建的对应模型)
更新:
显然在转换 mobilenet v2 模型后,tensorflow 冻结图包含比原始 pytorch 模型(~38 000 vs ~180)更多的卷积操作,如本 github issue中所述。
推荐阅读
- azure - 为什么使用 Azure Active Directory 调用我的 WebAPI 的 SPA 收到“此请求的授权已被拒绝。”?
- android - 显示来自这个特定 url 的图像
- sqlite - 如何在 ChangeNotifier 中使用 Futures?
- java - 无法从 Firebase 下载图像、文本、视频
- powerbi - 基于另一列的条件列
- c# - RemoveRange 删除的行数多于过滤后的行数
- docker - 安装 ElasticSearch 时,efk_elasticsearch_1 以代码 78 退出
- flutter - Flutter Form currentState on multiple tabbar NoSuchMethodError: The getter 'currentState' was called on null
- java - Curator Framework 绕过 Zookeeper ACL 设置
- android - 创建一个带有圆圈文本的自定义按钮