TFLite running in Cloud TPUs

General Discussion
,
#1

Hi!

Is it possible to run predictions using a TFLite quantized model in Cloud TPU Nodes?

I have already used a “traditional” Tensorflow model in TPU Nodes in order to run a object detection model. It ran perfectly.

The “traditional” version:

# Initilizing TPU System
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_name,
    zone=resource_zone,
    project=project_name
)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)

with tf.distribute.TPUStrategy(tpu_cluster_resolver):
    model = MyNNModel(...) # This is a Keras functional model
    model.load_weights('gs://my_bucket/path/to/checkpoints') 

results = model.predict(tfrecord_reference)

Now, using a TFLite model:

# Initilizing TPU System
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_name,
    zone=resource_zone,
    project=project_name
)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)

interpreter = tf.lite.Interpreter('gs://my_bucket/path/to/model.tflite')
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

outputs = interpreter.set_tensor(input_details[0], tfrecord_reference)
interpreter.invoke()
results = interpreter.get_tensor(output_details[0])

I receaved the message error message:

 interpreter = tf.lite.Interpreter('gs://my_bucket/path/to/model.tflite')
File "/usr/local/lib/python3.8/site-packages/tensorflow/lite/python/interpreter.py", line 365, in __init__
    _interpreter_wrapper.CreateWrapperFromFile(
ValueError: Could not open 'gs://my_bucket/path/to/model.tflite'.

I’ve tested the same TFLite model running a prediction in my Laptop and it worked correctly

The setup:
* TPU Node v2-8 (not preemptible)
* VM Instance running Ubuntu 21.04
* Tensorflow version 2.6.1
* Python 3.8
* TFLite model converted from a “tradicional” one
* TFLite model file uploaded to a GCP bucket
* VM Instance, TPU Node and GCP Bucket located at the same region

Thanks

#3

I don’t think it is currently possible to run TFLite models directly on Cloud TPUs. TensorFlow Lite is designed for on-device inference (such as mobile or edge devices), and Cloud TPUs require models in the standard TensorFlow format.