TensorFlow Lite on Unity

Style Transfer

TensorFlow Lite is a lightweight inference engine released as a sub package of TensorFlow. It runs on mobile and supports GPU acceleration with Metal on iOS/macOS, OpenGL on Android. A simple Unity plugin (the only single add ops network model) has been added in the experiments directory.

Unity already has developed official ML libraries called ML-Agent and Barracuda as an inference engine. It supports the ONNX format, so you can convert the model from TensorFlow or PyTorch... But I thought it would be great if I could use the pre-trained TFLite model without converting to the ONNX format. That’s why I started to test some models on TFLite examples to Unity, making some pull-requests to improve the Unity plugin which includes GPU support on Unity. Here is my repository.

asus4/tf-lite-unity-sample

How to use

Using the TFLite is simple.

1. Prepare the input array
2. Run `interpreter.Invoke();`
3. Postprocess the output array

Here is a simple MNIST script.

Simple usage of MNIST

GPU delegate option will dramatically improve the speed if your model has image inputs/outputs. At the moment, it does not support passing GPU-texture as models’input directly, despite Unity Barracuda supports it.

How to use Metal Delegate option

Usually, the model using Image needs pre/post-processing. Unity makes it easier and faster than native implementations.

The resize / crop sample

ShowCases

SSD Object Detection

SSD (single shot detector) Object Detection.

DeepLab Semantic Segmentation

DeepLab Semantic Segmentation.

Text classification

Also, the Text Classification worked on the CPU-mode.

Style Transfer

A practical sample that uses Style Transfer as a camera filter.

Future development plans

• Support direct passing Unity texture to TensorFlow GPU delegate to reduce the exchanging textures between the GPU and CPU.
• Support custom Ops insertion to use MediaPipe examples.