❗ This codebase has been archived because there are more updated solutions available. Please check out https://ai.google.dev/edge/lite/models/pytorch_to_tflite

PyTorch to TensorFlow Lite Converter

Converts PyTorch whole model into Tensorflow Lite

PyTorch -> Onnx -> Tensorflow 2 -> TFLite

Please install first

python3 setup.py install

Args

• --torch-path Path to local PyTorch model, please save whole model e.g. torch.save(model, PATH)
• --tf-lite-path Save path for Tensorflow Lite model
• --target-shape Model input shape to create static-graph (default: (224, 224, 3)
• --sample-file Path to sample image file. If model is not about computer-vision, please use leave empty and only enter --target-shape
• --seed Seeds RNG to produce random input data when --sample-file does not exists
• --log=INFO To see what happens behind

Basic usage of the script

To test with sample file:

python3 -m torch2tflite.converter
    --torch-path tests/mobilenetv2_model.pt
    --tflite-path mobilenetv2.tflite
    --sample-file sample_image.png
    --target-shape 224 224 3

To test with random input to check gradients:

python3 -m torch2tflite.converter
    --torch-path tests/mobilenetv2_model.pt
    --tflite-path mobilenetv2.tflite
    --target-shape 224 224 3
    --seed 10