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

Pytorch:将2D-CNN模型转换为tflite,pytorch,onnx,tensorflow-lite,Pytorch,Onnx,Tensorflow Lite,我想将一个模型(如Mobilenet V2)从pytorch转换为tflite,以便在移动设备上运行它 有人做到了吗 我发现的只是一种使用ONNX将模型转换为中间状态的方法。然而,这似乎不能正常工作,因为Tensorflow期望NHWC通道顺序,而onnx和pytorch则使用NCHW通道顺序 有一个讨论,但在我的案例中,转换工作没有抱怨,直到“冻结的张量流图模型”,在尝试将模型进一步转换为tflite后,它抱怨通道顺序错误 以下是我目前的代码: import torch import torc

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

有人做到了吗

我发现的只是一种使用ONNX将模型转换为中间状态的方法。然而,这似乎不能正常工作,因为Tensorflow期望NHWC通道顺序,而onnx和pytorch则使用NCHW通道顺序

有一个讨论,但在我的案例中,转换工作没有抱怨,直到“冻结的张量流图模型”,在尝试将模型进一步转换为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警告)。然后我使用(“input.1”和“473”)查找输入和输出张量的名称

最后,我将我常用的tf图应用于来自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 
torch                1.6.0.dev20200508 
tensorflow-gpu       2.2.0
onnx                 1.7.0              
onnx-tf              1.5.0 
下面是我从
tflite
得到的确切错误消息:

Unexpected value for attribute 'data_format'. Expected 'NHWC'
Fatal Python error: Aborted
更新
正在更新我的配置:

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 
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
tflite\u转换\
--输出文件=mobilenet\u v2.tflite\
--graph_def_file=mobilenet_v2.pb\
--input_数组=input.1\
--输出_阵列=473\
--启用_v1_converter#@Ahwar使用谷歌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对180)多得多,如中所述。

您可能需要升级您的tensorflow版本,1.14使用的旧转换器不支持2.2那样多的型号。升级到tensorflow 2.2会导致另一个错误,而转换到tflite:
:0:error:loc(“卷积”):“tf.Conv2D”op既不是自定义op,也不是flex op
。我想我需要兼容版本的精确组合。这是一个完全依赖的地狱…很抱歉的挫折感-这应该是可行的,但很难说不知道pb中有什么。你能在某处发布Netron或graphdef的截图吗?@daverim我添加了Netron的图片和模型的链接(正如我所说:这些是“未接触过的”mobilenet v2模型,所以我想它们至少应该与一些配置配合使用,@mcExchange支持我的答案和传播。