tensorflow - How do you convert a .onnx to tflite?

Question

I've exported my model to ONNX via:

# Export the model
torch_out = torch.onnx._export(learn.model,             # model being run
                           x,                       # model input (or a tuple for multiple inputs)
                          EXPORT_PATH + "mnist.onnx", # where to save the model (can be a file or file-like object)
                           export_params=True)      # store the trained parameter weights inside the model file

And now I am trying to convert the model to a Tensorflow Lite file so that I can do inference on Android. Unfortunately, PyTorch/Caffe2 support is fairly lacking or too complex for Android but Tensorflow appears much simpler.

The documentation for ONNX to Tflite is pretty light on this.

I've tried exporting to a Tensorflow GraphDef proto via:

tf_rep.export_graph(EXPORT_PATH + 'mnist-test/mnist-tf-export.pb')

And then running toco:

toco 
--graph_def_file=mnist-tf-export.pb 
--input_format=TENSORFLOW_GRAPHDEF 
--output_format=TFLITE 
--inference_type=FLOAT 
--input_type=FLOAT 
--input_arrays=0 
--output_arrays=add_10 
--input_shapes=1,3,28,28 
--output_file=mnist.tflite`

When I do though I get the following error:

File "anaconda3/lib/python3.6/site-packages/tensorflow/lite/python/convert.py", line 172, in toco_convert_protos
    "TOCO failed. See console for info.
%s
%s
" % (stdout, stderr))
tensorflow.lite.python.convert.ConverterError: TOCO failed. See console for info.
2018-11-06 16:28:33.864889: I tensorflow/lite/toco/import_tensorflow.cc:1268] Converting unsupported operation: PyFunc
2018-11-06 16:28:33.874130: F tensorflow/lite/toco/import_tensorflow.cc:114] Check failed: attr.value_case() == AttrValue::kType (1 vs. 6)

Further, even when I run the command I don't know what to specify for the input_arrays or output_arrays since the model was originally built in PyTorch.

Has anyone successfully converted their ONNX model to TFlite?

Here's the ONNX file I'm trying to convert: https://drive.google.com/file/d/1sM4RpeBVqPNw1WeCROpKLdzbSJPWSK79/view?usp=sharing

Extra info

• Python 3.6.6 :: Anaconda custom (64-bit)
• onnx.version = '1.3.0'
• tf.version = '1.13.0-dev20181106'
• torch.version = '1.0.0.dev20181029'

Answer 1

I think the ONNX file i.e. model.onnx that you have given is corrupted I don't know what is the issue but it is not doing any inference on ONNX runtime.

Now you can run PyTorch Models directly on mobile phones. check out PyTorch Mobile's documentation here

This answer is for TensorFlow version 1,
For TensorFlow version 2 or higher click link

The best way to convert the model from protobuf freezeGraph to TFlite is to use the official TensorFlow lite converter documentation

According to TensorFlow Docs, TocoConverter has been deprecated

This class (tf.compat.v1.lite.TocoConverter) has been deprecated. Please use lite.TFLiteConverter instead.

Convert from PyTorch to ONNX model

The best practice to convert the model from Pytorch to Onnx is that you should add the following parameters to specify the names of the input and output layer of your model in torch.onnx.export() function

# Export the model from PyTorch to ONNX
torch_out = torch.onnx._export(model,             # model being run
                                x,          # model input (or a tuple for multiple inputs)
                                EXPORT_PATH + "mnist.onnx",      # where to save the model (can be a file or file-like object)
                                export_params=True,       # store the trained parameter weights inside the model file
                                input_names=['main_input'],     # specify the name of input layer in onnx model
                                output_names=['main_output'])     # specify the name of input layer in onnx model

So in your case: Now export this model to TensorFlow protobuf FreezeGraph using onnx-tf

Please note that this method is only working when tensorflow_version < 2

Convert from ONNX to TensorFlow freezGraph

To convert the model please install onnx-tf version 1.5.0 from the below command

pip install  onnx-tf==1.5.0

Now to convert .onnx model to TensorFlow freeze graph run this below command in shell

onnx-tf convert -i "mnist.onnx" -o  "mnist.pb"

Convert from TensorFlow FreezeGraph .pb to TF

Now to convert this model from .pb file to tflite model use this code

import tensorflow as tf
# make a converter object from the saved tensorflow file
converter = tf.lite.TFLiteConverter.from_frozen_graph('mnist.pb', #TensorFlow freezegraph .pb model file
                                                      input_arrays=['main_input'], # name of input arrays as defined in torch.onnx.export function before.
                                                      output_arrays=['main_output']  # name of output arrays defined in torch.onnx.export function before.
                                                      )
# tell converter which type of optimization techniques to use
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# to view the best option for optimization read documentation of tflite about optimization
# go to this link https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional

# convert the model 
tf_lite_model = converter.convert()
# save the converted model 
open('mnist.tflite', 'wb').write(tf_lite_model)

To choose which option is best for optimization for your model use case see this official guide about TensorFlow lite optimization

https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional

Note: You can try my Jupyter Notebook Convert ONNX model to Tensorflow Lite on Google Colaboratory link