Run TensorFlow models

TensorFlow is an open-source machine learning framework developed by Google that provides tools for building, training, and deploying neural networks. To run TensorFlow models on the NPU of your Dragonwing development board you'll need to convert your model to a quantized TFLite model. You can then use LiteRT to run the model with full hardware acceleration.

Quantizing and converting a model

Your TensorFlow model currently uses 32-bit floating point numbers for its weights and activations. The NPU on your development board only supports 8-bit integers, so the model must be quantized - converted from floating point to fixed point values. This makes the model smaller and faster to run (and run on the NPU), but it has an effect on accuracy.

The easiest way to quantize models is using the post-training quantization. Here you take an already trained model, then quantize the weights and activations using a representative dataset. This means that there's no effect in the training loop. Optionally, you can also add TensorFlow's built-in Quantization aware training to reduce quantization error (but it requires changes to your training loop).

Let's demonstrate by quantizing a Keras model. Open the terminal on your development board, or an ssh session to your development board, and:

  1. Create a new venv and install some base packages:

    mkdir -p ~/post-training-quantization-tf
cd ~/post-training-quantization-tf

python3 -m venv .venv
source .venv/bin/activate

pip3 install tensorflow==2.20.0 tf_keras==2.20.1 ai-edge-litert==1.3.0

  2. Download a demonstration model in .keras format, plus a test set.

    mkdir -p models
wget -O models/cats.keras https://cdn.edgeimpulse.com/qc-ai-docs/models/cats.keras
wget -O models/cats_X_val.npy https://cdn.edgeimpulse.com/qc-ai-docs/models/cats_X_val.npy
wget -O models/cats_y_val.npy https://cdn.edgeimpulse.com/qc-ai-docs/models/cats_y_val.npy

  3. Create a new file quantize.py and add:

    import tensorflow as tf, numpy as np, os, time, tf_keras as keras
from ai_edge_litert.interpreter import Interpreter, load_delegate

# Shape: (444, 160, 160, 3)
X_val = np.load('models/cats_X_val.npy')
# Shape: (444, 1) -> with class 1..6 -> scale to 0..5
y_val = np.load('models/cats_y_val.npy') - 1

# Load Keras model
model = keras.models.load_model("models/cats.keras")

# Calculate accuracy of the TF model
tf_start = time.perf_counter()
y_pred = model.predict(X_val)
tf_end = time.perf_counter()
preds = np.argmax(y_pred, axis=1)
acc_tf = (preds == y_val).mean()
print(f"TF/Keras accuracy: {acc_tf*100:.2f}% (time per inference: {(tf_end - tf_start) * 1000 / X_val.shape[0]:.4g}ms)")
print('')

# Convert to quantized TFLite file... Uses the dataset earlier as a representative dataset to improve accuracy.
TFLITE_FILE = 'cats_i8.tflite'
if not os.path.exists(TFLITE_FILE):
    print(f'Converting to TFLite file ({TFLITE_FILE})...')

    def rep_dataset():
        for i in range(X_val.shape[0]):
            yield [X_val[i:i+1]]

    # Build a fixed batch=1 input signature (QNN cannot handle dynamic dims)
    specs = []
    for t in model.inputs:
        if None in t.shape[1:]:
            raise ValueError(f"Non-batch dims must be known; got {t.shape}")
        specs.append(tf.TensorSpec([1, *t.shape[1:]], dtype=t.dtype, name=t.name.split(':')[0]))

    @tf.function(input_signature=specs)
    def serve(*xs):
        y = model(*xs)
        return y if isinstance(y, (tuple, list)) else (y,)  # keep output order stable

    concrete = serve.get_concrete_function()
    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete], model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = rep_dataset
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.inference_input_type = tf.int8
    converter.inference_output_type = tf.int8
    tflite_model = converter.convert()
    with open(TFLITE_FILE, "wb") as f:
        f.write(tflite_model)

    print(f"TFLite written: {TFLITE_FILE} ({os.path.getsize(TFLITE_FILE)/1e6:.2f} MB)")
else:
    print(f'TFLite file already exists ({TFLITE_FILE})')
print('')

def run_tflite_model(model_path, use_npu):
    # Use QNN to run this model on NPU
    experimental_delegates = []
    if use_npu:
        experimental_delegates = [load_delegate("libQnnTFLiteDelegate.so", options={"backend_type": "htp"})]

    # Get accuracy for the quantized TFLite file, construct the interpreter
    interpreter = Interpreter(model_path=model_path, experimental_delegates=experimental_delegates)
    interpreter.allocate_tensors()
    in_details = interpreter.get_input_details()[0]
    out_details = interpreter.get_output_details()[0]

    # You need to scale the input / output yourself using quantization params
    in_scale, in_zp = in_details["quantization"]
    out_scale, out_zp = out_details["quantization"]

    # Loop through one-by-one (most TFLite files have a fixed batch size of 1)
    preds_tflite = []
    tflite_start = time.perf_counter()
    for i in range(X_val.shape[0]):
        # Scale input and invoke
        x = X_val[i:i+1]
        x_q = np.round(x / in_scale + in_zp).astype(in_details['dtype'])
        interpreter.set_tensor(in_details["index"], x_q)
        interpreter.invoke()
        # Scale output back to f32
        out = interpreter.get_tensor(out_details["index"])
        out = (out.astype(np.float32) - out_zp) * out_scale
        # And add the outcome to the predictions
        preds_tflite.append(np.argmax(out, axis=1)[0])
    tflite_end = time.perf_counter()

    # Compare accuracy in the same way as above
    acc_tflite = (np.array(preds_tflite) == y_val).mean()
    if use_npu:
        print(f"Quantized TFLite accuracy (NPU): {acc_tflite*100:.2f}% (time per inference: {(tflite_end - tflite_start) * 1000 / X_val.shape[0]:.4g}ms)")
    else:
        print(f"Quantized TFLite accuracy (CPU): {acc_tflite*100:.2f}% (time per inference: {(tflite_end - tflite_start) * 1000 / X_val.shape[0]:.4g}ms)")

run_tflite_model(TFLITE_FILE, False)
run_tflite_model(TFLITE_FILE, True)

  4. Run the example:

    python3 quantize.py

# TF/Keras accuracy: 94.37% (time per inference: 29.86ms)
#
# Converting to TFLite file (cats_i8.tflite)...
# ...
# Quantized TFLite accuracy (CPU): 87.16% (time per inference: 10.37ms)
# Quantized TFLite accuracy (NPU): 88.51% (time per inference: 3.809ms)

Great! You now have cats_i8.tflite which runs ~8x faster on the NPU (but with some accuracy loss). See Run LiteRT/TFLite models on NPU for more details on the LiteRT runtime (including C++ examples).

PreviousQualcomm AI HubNextRun LiteRT / TFLite models

Last updated