Export and validate

import json, numpy as np, timm, torch
from pathlib import Path
from torchvision import datasets, transforms as T

ROOT = Path.home() / "sutraai"; EXP = ROOT / "exports"; EXP.mkdir(exist_ok=True)
SIZE, mean, std = 96, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
eval_tf = T.Compose([T.Resize(int(SIZE * 1.15)), T.CenterCrop(SIZE), T.ToTensor(), T.Normalize(mean, std)])
raw_tf = T.Compose([T.Resize(int(SIZE * 1.15)), T.CenterCrop(SIZE), T.ToTensor()])        # 0–1, no normalize
test_ds = datasets.ImageFolder(ROOT / "data/processed/images/test", eval_tf)
classes = test_ds.classes

model = timm.create_model("mobilenetv2_050", pretrained=False, num_classes=len(classes))
model.load_state_dict(torch.load(ROOT / "models/image_v1/best.pt", map_location="cpu"))
model.eval()
X = torch.stack([test_ds[i][0] for i in range(len(test_ds))]); y = np.array(test_ds.targets)
with torch.no_grad(): float_acc = float((model(X).argmax(1).numpy() == y).mean())
print("float32 test accuracy:", round(float_acc, 3))

import onnxruntime as ort, warnings
warnings.filterwarnings("ignore", message="You are using the legacy TorchScript-based ONNX export")   # expected on torch 2.8, see below
onnx_path = EXP / "image_v1.onnx"
torch.onnx.export(model, torch.randn(1, 3, SIZE, SIZE), str(onnx_path), dynamo=False,   # classic exporter: simplest graph for TensorRT + onnx2tf
                  input_names=["input"], output_names=["logits"], opset_version=17)
# The model is exported with a fixed batch size of 1 (what TensorRT and microcontrollers use),
# so the check below feeds it one image at a time.

sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])   # CPU = exact float32, a fair comparison
n = min(20, len(X))
with torch.no_grad(): ref = model(X[:n]).numpy()
out = np.concatenate([sess.run(None, {"input": X[i:i + 1].numpy()})[0] for i in range(n)])
onnx_diff = float(np.abs(ref - out).max())
print("max difference PyTorch vs ONNX:", onnx_diff)         # expect: below 0.001, e.g. 2e-05 (= 0.00002) is a pass

!/usr/src/tensorrt/bin/trtexec --onnx={onnx_path} --saveEngine={EXP}/image_v1_fp16.engine --fp16 2>&1 | grep -E "Latency|Throughput|PASSED|FAILED"

!/usr/src/tensorrt/bin/trtexec --onnx={onnx_path} --saveEngine={EXP}/image_v1_dla0.engine --fp16 --useDLACore=0 --allowGPUFallback 2>&1 | grep -E "Latency|DLA|PASSED|FAILED"

calib_ds = datasets.ImageFolder(ROOT / "data/processed/images/train", raw_tf)
idx = np.random.default_rng(0).choice(len(calib_ds), size=min(100, len(calib_ds)), replace=False)
calib = np.stack([calib_ds[i][0].permute(1, 2, 0).numpy() for i in idx]).astype(np.float32)   # N,H,W,C in 0–1
np.save(EXP / "calib_image.npy", calib)
print(calib.shape)                     # e.g. (100, 96, 96, 3)

mkdir C:\sutraai-export -Force
scp <your-user>@<jetson-ip>:~/sutraai/exports/image_v1.onnx <your-user>@<jetson-ip>:~/sutraai/exports/calib_image.npy C:\sutraai-export\

!pip install -q onnx2tf==2.6.9
from google.colab import files
files.upload()           # choose image_v1.onnx and calib_image.npy from C:\sutraai-export

!onnx2tf -i image_v1.onnx -o image_v1_tf -oiqt -qt per-tensor -cind input calib_image.npy "[[[[0.485,0.456,0.406]]]]" "[[[[0.229,0.224,0.225]]]]"
# it may end with an error about "int16_act": fine, as long as the next line lists image_v1_full_integer_quant.tflite
!ls image_v1_tf/*.tflite
files.download("image_v1_tf/image_v1_full_integer_quant.tflite")

ssh <your-user>@<jetson-ip> "mkdir -p ~/sutraai/exports/image_v1_tf"
scp $HOME\Downloads\image_v1_full_integer_quant.tflite <your-user>@<jetson-ip>:~/sutraai/exports/image_v1_tf/

sudo apt-get install -y cmake build-essential                 # needed to compile onnxoptimizer and psutil
curl -LsSf https://astral.sh/uv/install.sh | sh                # installs uv into ~/.local/bin
cd ~/sutraai
~/.local/bin/uv venv --python 3.12 .venv-convert               # downloads a standalone Python 3.12 just for this venv
~/.local/bin/uv pip install --python .venv-convert/bin/python onnx2tf==2.6.9
.venv-convert/bin/onnx2tf -h 2>/dev/null | head -n 3           # converter works: prints "usage: onnx2tf [-h] ..."

import os, shutil, subprocess
from pathlib import Path

CONVERTER = Path.home() / "sutraai/.venv-convert/bin"          # Python 3.12 venv with onnx2tf 2.6.9 (Step 05)

def onnx2tf_int8(onnx_path, out_dir, calib_npy, mean, std, per_tensor=True, backend=None, converter=CONVERTER):
    """ONNX → full-integer int8 TFLite. Returns the path of <name>_full_integer_quant.tflite.

    per_tensor=True works around an onnx2tf 2.6.9 bug: per-channel quantization of depthwise
    convolutions (MobileNetV2) fails validation. onnx2tf may still exit with an error while writing
    the optional int16-activation variants; that's fine as long as the full-integer file exists."""
    onnx_path, out_dir, converter = Path(onnx_path), Path(out_dir), Path(converter)
    shutil.rmtree(out_dir, ignore_errors=True)                     # never pick up a stale file
    env = {**os.environ, "PATH": f"{converter}:{os.environ.get('PATH', '')}"}   # onnx2tf runs the onnxsim tool by name
    cmd = [str(converter / "onnx2tf"), "-i", str(onnx_path), "-o", str(out_dir), "-oiqt",
           "-cind", "input", str(calib_npy), mean, std]
    if per_tensor: cmd += ["-qt", "per-tensor"]
    if backend: cmd += ["-tb", backend]
    r = subprocess.run(cmd, env=env, capture_output=True, text=True)
    tfl = out_dir / f"{onnx_path.stem}_full_integer_quant.tflite"
    if not tfl.exists():
        print(r.stdout[-3000:], r.stderr[-3000:])
        raise RuntimeError(f"onnx2tf did not create {tfl.name} (exit code {r.returncode})")
    return tfl

import sys; sys.path.insert(0, str(ROOT))
from sutra_core.export import onnx2tf_int8
tfl = onnx2tf_int8(onnx_path, EXP / "image_v1_tf", EXP / "calib_image.npy",
                   "[[[[0.485,0.456,0.406]]]]", "[[[[0.229,0.224,0.225]]]]")
print(tfl.name, round(tfl.stat().st_size / 1024, 1), "KB")      # expect: image_v1_full_integer_quant.tflite, a few hundred KB

from ai_edge_litert.interpreter import Interpreter

def tflite_accuracy(path, X_nhwc, y):
    """Full-integer model. X_nhwc: float inputs already normalized, in N,H,W,C order."""
    itp = Interpreter(model_path=str(path)); itp.allocate_tensors()
    inp, out = itp.get_input_details()[0], itp.get_output_details()[0]
    scale, zero = inp["quantization"]
    hits = 0
    for xi, yi in zip(X_nhwc, y):
        q = np.clip(np.round(xi / scale + zero), -128, 127).astype(inp["dtype"])
        itp.set_tensor(inp["index"], q[None]); itp.invoke()
        hits += int(itp.get_tensor(out["index"])[0].argmax() == yi)
    return hits / len(y)

tfl = EXP / "image_v1_tf" / "image_v1_full_integer_quant.tflite"
int8_acc = tflite_accuracy(tfl, X.permute(0, 2, 3, 1).numpy(), y)
print(f"float32 {float_acc:.3f}  →  int8 {int8_acc:.3f}   (drop {100 * (float_acc - int8_acc):.1f} points)")
print("int8 model size:", round(tfl.stat().st_size / 1024, 1), "KB")

!cd {EXP} && xxd -i image_v1_tf/image_v1_full_integer_quant.tflite > image_v1_model_data.cc && head -c 300 image_v1_model_data.cc