PyTorch to TFLite flow

DeepGate's native route is dg lowering a model straight to schema.json (see Export model). The compiler also accepts a fully int8-quantized TensorFlow Lite (.tflite) file through a beta front-end. There is no single call that gets you there: litert-torch converts your model to a float .tflite, then ai-edge-quantizer quantizes that file to int8.

pip install litert-torch ai-edge-quantizer     # litert-torch was formerly ai-edge-torch

import litert_torch
import numpy as np
import torch
from ai_edge_litert.interpreter import Interpreter
from ai_edge_quantizer import quantizer, recipe

# Convert the float model (model: your trained nn.Module, input [1, 1, 28, 28])
sample = (torch.randn(1, 1, 28, 28),)          # batch must be 1 for TFLite
litert_torch.convert(model.eval(), sample).export("float.tflite")

# Build calibration data, keyed by the float model's signature and input names
sig = Interpreter(model_path="float.tflite").get_signature_list()
sig_key = next(iter(sig))
input_name = sig[sig_key]["inputs"][0]
samples = [np.random.randn(1, 1, 28, 28).astype(np.float32) for _ in range(32)]

# Quantize the file to full-integer int8
qt = quantizer.Quantizer("float.tflite")
qt.load_quantization_recipe(recipe.static_wi8_ai8())
result = qt.calibrate({sig_key: [{input_name: x} for x in samples]})
qt.quantize(result).export_model("int8.tflite")

Calibrate on real representative inputs, not the random ones above, or the activation ranges will be wrong.

Keep recipe.static_wi8_ai8(): it quantizes both the weights and the activations to int8, which is what the compiler needs. Recipes that leave any part of the model in float will not compile.