Pytorch:将 2D-CNN 模型转换为 tflite

Question

我想将模型(例如 Mobilenet V2)从 pytorch 转换为 tflite，以便在移动设备上运行它。

有人成功了吗？

我发现的是一种使用 ONNX 将模型转换为中间状态的方法。然而，这似乎无法正常工作，因为 Tensorflow 需要 NHWC channel 顺序，而 onnx 和 pytorch 使用 NCHW channel 顺序。

有讨论on github ，但是就我而言，转换工作没有任何提示，直到“卡住 tensorflow 图模型”，在尝试将模型进一步转换为 tflite 后，它提示 channel 顺序错误......

这是迄今为止我的代码:

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 模型:

Here是指向我转换后的 onnx 和 pb 文件的 gdrive 链接

Answer 1

@Ahwar 发布了一个不错的 solution使用 Google Colab 笔记本来完成此操作。

它使用

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 )，如 this github issue 中所述。 .