Pytorch: Convert 2D-CNN model to tflite
I'd like to convert a model (eg Mobilenet V2) from pytorch to tflite in order to run it on a mobile device.

Has anyone managed to do so?

All I found, was a method that uses ONNX to convert the model into an inbetween state. However, this seems not to work properly, as Tensorflow expects a NHWC-channel order whereas onnx and pytorch work with NCHW channel order.

There is a discussion on github, however in my case the conversion worked without complaints until a "frozen tensorflow graph model", after trying to convert the model further to tflite, it complains about the channel order being wrong...

Here is my code so far:

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) 


All working without errors until here (ignoring many tf warnings). Then I look up the names of the input and output tensors using netron ("input.1" and "473").

Finally I apply my usual tf-graph to tf-lite conversion script from bash:

tflite_convert \
    --output_file=mobilenet_v2.tflite \
    --graph_def_file=mobilenet_v2.pb \
    --input_arrays=input.1 \
    --output_arrays=473


My configuration:

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 


Here is the exact error message I'm getting from tflite:

Unexpected value for attribute 'data_format'. Expected 'NHWC'
Fatal Python error: Aborted


UPDATE:
Updating my configuration:

torch                1.6.0.dev20200508 
tensorflow-gpu       2.2.0
onnx                 1.7.0              
onnx-tf              1.5.0 


using 

tflite_convert \
    --output_file=mobilenet_v2.tflite \
    --graph_def_file=mobilenet_v2.pb \
    --input_arrays=input.1 \
    --output_arrays=473 \
    --enable_v1_converter  # &lt;-- needed for conversion of frozen graphs


leading to another error:

Exception: &lt;unknown&gt;:0: error: loc("convolution"): 'tf.Conv2D' op is neither a custom op nor a flex op


Update:
Here is an onnx model of mobilenet v2 loaded via netron:  



Here  is a gdrive link to my converted onnx and pb file

@Ahwar posted a nice solution to this using a Google Colab notebook.

It uses

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


The conversion is working and the model can be tested on my computer. However when pushing the model to the mobile phone it only works in CPU mode and is much slower (almost 10 fold) than a corresponding model created in tensorflow directly. GPU mode is not working on my mobile phone (in contrast to the corresponding model created in tensorflow directly)

Update:
Apparantly after converting the mobilenet v2 model, the tensorflow frozen graph contains many more convolution operations than the original pytorch model ( ~38 000 vs ~180 ) as discussed in this github issue.
