Pytorch: Convert 2D-CNN model to tflite

Question

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  # <-- needed for conversion of frozen graphs

leading to another error:

Exception: <unknown>: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

Answer 1

@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.

Answer 2

https://github.com/alibaba/TinyNeuralNetwork

You can try this project to convert the pytorch model to tflite. It supports all models in torchvision, and can eliminate redundant operators, basically without performance loss