Question

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

有人能这样做吗？

我发现的所有方法都是使用ONNX将模型转换为中间状态的方法。但是，这似乎无法正常工作，因为Tensorflow期望使用NHWC通道顺序，而onnx和pytorch可以使用NCHW通道顺序。

有一个讨论on 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模型：

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

Answer 1

@Ahwar使用Google Colab笔记本发布了一个不错的solution。

它使用

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模型后，张量流冻结图包含的卷积操作要比this github issue中讨论的原始pytorch模型（〜38 000 vs〜180）更多。

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

1 个答案: