export_litert.py

#!/usr/bin/env python3
"""
Convert InsectCNN PyTorch weights into a Google LiteRT / TensorFlow Lite (.tflite) model.
Validates numerical parity between PyTorch and LiteRT outputs.
"""

import argparse
import json
from pathlib import Path
import numpy as np
import torch
import tensorflow as tf
from src.inference import load_model


class InsectLiteRT(tf.Module):
    """Exact TensorFlow representation of InsectCNN for LiteRT export."""
    def __init__(self, pt_net):
        super().__init__()
        sd = pt_net.state_dict()

        # Block 1: Conv2d(1, 16, 3, pad 1) + BatchNorm2d(16)
        # PyTorch Conv2d weight: (out_channels, in_channels, kH, kW)
        # TF Conv2D weight: (kH, kW, in_channels, out_channels)
        self.w0 = tf.Variable(sd['features.0.weight'].cpu().numpy().transpose(2, 3, 1, 0), name='conv0_w')
        self.b0 = tf.Variable(sd['features.0.bias'].cpu().numpy(), name='conv0_b')
        self.bn0_gamma = tf.Variable(sd['features.1.weight'].cpu().numpy(), name='bn0_gamma')
        self.bn0_beta = tf.Variable(sd['features.1.bias'].cpu().numpy(), name='bn0_beta')
        self.bn0_mean = tf.Variable(sd['features.1.running_mean'].cpu().numpy(), name='bn0_mean')
        self.bn0_var = tf.Variable(sd['features.1.running_var'].cpu().numpy(), name='bn0_var')
        self.bn0_eps = tf.constant(1e-5, dtype=tf.float32)

        # Block 2: Conv2d(16, 32, 3, pad 1) + BatchNorm2d(32)
        self.w4 = tf.Variable(sd['features.4.weight'].cpu().numpy().transpose(2, 3, 1, 0), name='conv4_w')
        self.b4 = tf.Variable(sd['features.4.bias'].cpu().numpy(), name='conv4_b')
        self.bn4_gamma = tf.Variable(sd['features.5.weight'].cpu().numpy(), name='bn4_gamma')
        self.bn4_beta = tf.Variable(sd['features.5.bias'].cpu().numpy(), name='bn4_beta')
        self.bn4_mean = tf.Variable(sd['features.5.running_mean'].cpu().numpy(), name='bn4_mean')
        self.bn4_var = tf.Variable(sd['features.5.running_var'].cpu().numpy(), name='bn4_var')
        self.bn4_eps = tf.constant(1e-5, dtype=tf.float32)

        # Block 3: Conv2d(32, 64, 3, pad 1) + BatchNorm2d(64)
        self.w8 = tf.Variable(sd['features.8.weight'].cpu().numpy().transpose(2, 3, 1, 0), name='conv8_w')
        self.b8 = tf.Variable(sd['features.8.bias'].cpu().numpy(), name='conv8_b')
        self.bn8_gamma = tf.Variable(sd['features.9.weight'].cpu().numpy(), name='bn8_gamma')
        self.bn8_beta = tf.Variable(sd['features.9.bias'].cpu().numpy(), name='bn8_beta')
        self.bn8_mean = tf.Variable(sd['features.9.running_mean'].cpu().numpy(), name='bn8_mean')
        self.bn8_var = tf.Variable(sd['features.9.running_var'].cpu().numpy(), name='bn8_var')
        self.bn8_eps = tf.constant(1e-5, dtype=tf.float32)

        # Classifier Linear 1: (1024 -> 128)
        self.fc2_w = tf.Variable(sd['classifier.2.weight'].cpu().numpy().T, name='fc2_w')
        self.fc2_b = tf.Variable(sd['classifier.2.bias'].cpu().numpy(), name='fc2_b')

        # Classifier Linear 2: (128 -> num_classes)
        self.fc5_w = tf.Variable(sd['classifier.5.weight'].cpu().numpy().T, name='fc5_w')
        self.fc5_b = tf.Variable(sd['classifier.5.bias'].cpu().numpy(), name='fc5_b')

    @tf.function(input_signature=[tf.TensorSpec(shape=[1, 1, 64, 201], dtype=tf.float32, name='audio_mel')])
    def __call__(self, x):
        # Input x is NCHW: (1, 1, 64, 201) -> convert to NHWC: (1, 64, 201, 1)
        x = tf.transpose(x, [0, 2, 3, 1])

        # Block 1
        x = tf.nn.conv2d(x, self.w0, strides=1, padding='SAME') + self.b0
        x = (x - self.bn0_mean) / tf.sqrt(self.bn0_var + self.bn0_eps) * self.bn0_gamma + self.bn0_beta
        x = tf.nn.relu(x)
        x = tf.nn.max_pool2d(x, ksize=2, strides=2, padding='VALID')

        # Block 2
        x = tf.nn.conv2d(x, self.w4, strides=1, padding='SAME') + self.b4
        x = (x - self.bn4_mean) / tf.sqrt(self.bn4_var + self.bn4_eps) * self.bn4_gamma + self.bn4_beta
        x = tf.nn.relu(x)
        x = tf.nn.max_pool2d(x, ksize=2, strides=2, padding='VALID')

        # Block 3
        x = tf.nn.conv2d(x, self.w8, strides=1, padding='SAME') + self.b8
        x = (x - self.bn8_mean) / tf.sqrt(self.bn8_var + self.bn8_eps) * self.bn8_gamma + self.bn8_beta
        x = tf.nn.relu(x)

        # AdaptiveAvgPool2d((4, 4)) over feature map shape (16, 50)
        # Convert back to NCHW for matching PyTorch index ordering: (1, 64, 16, 50)
        x_nchw = tf.transpose(x, [0, 3, 1, 2])
        slices = []
        for oh in range(4):
            sh = oh * 4
            eh = (oh + 1) * 4
            for ow in range(4):
                sw = (ow * 50) // 4
                ew = ((ow + 1) * 50 + 3) // 4
                pool_slice = tf.reduce_mean(x_nchw[:, :, sh:eh, sw:ew], axis=[2, 3])
                slices.append(pool_slice)

        pooled = tf.stack(slices, axis=-1)  # (1, 64, 16)
        flat = tf.reshape(pooled, [1, 1024])

        # Dense classifier
        h = tf.matmul(flat, self.fc2_w) + self.fc2_b
        h = tf.nn.relu(h)
        logits = tf.matmul(h, self.fc5_w) + self.fc5_b
        return logits


def export_litert(model_dir='.', output_file='model.tflite', quantize=False):
    model_dir = Path(model_dir)
    output_path = Path(output_file)
    if not output_path.is_absolute():
        output_path = model_dir / output_path

    print(f"[*] Loading PyTorch model from '{model_dir}'...")
    pt_net, labels, cfg = load_model(model_dir)
    pt_net.eval()

    print("[*] Constructing TensorFlow graph with transferred weights...")
    tf_module = InsectLiteRT(pt_net)

    print(f"[*] Converting graph to LiteRT / TFLite format (quantize={quantize})...")
    concrete_func = tf_module.__call__.get_concrete_function()
    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func], trackable_obj=tf_module)
    if quantize:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
    tflite_bytes = converter.convert()

    output_path.write_bytes(tflite_bytes)
    print(f"[+] Saved LiteRT model to '{output_path}' ({len(tflite_bytes):,} bytes).")

    # Numerical verification
    print("[*] Running numerical parity test across 5 random inputs...")
    interpreter = tf.lite.Interpreter(model_content=tflite_bytes)
    interpreter.allocate_tensors()
    in_idx = interpreter.get_input_details()[0]['index']
    out_idx = interpreter.get_output_details()[0]['index']

    rng = np.random.default_rng(42)
    max_diff = 0.0
    for _ in range(5):
        sample = rng.standard_normal((1, 1, 64, 201)).astype(np.float32)
        with torch.no_grad():
            pt_out = pt_net(torch.from_numpy(sample)).cpu().numpy()

        interpreter.set_tensor(in_idx, sample)
        interpreter.invoke()
        tflite_out = interpreter.get_tensor(out_idx)

        diff = np.max(np.abs(pt_out - tflite_out))
        if diff > max_diff:
            max_diff = diff

    print(f"[+] Maximum absolute difference between PyTorch and LiteRT: {max_diff:.8e}")
    tol = 0.2 if quantize else 1e-4
    assert max_diff < tol, f"Parity check failed: max_diff {max_diff} exceeds {tol} tolerance!"
    print("[+] Numerical parity verified successfully!")
    return str(output_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Export PyTorch model to LiteRT / TFLite")
    parser.add_argument('--model-dir', default='.', help="Directory containing config.json and model.pt/model.h5")
    parser.add_argument('--output', default='model.tflite', help="Output filename for .tflite model")
    parser.add_argument('--quantize', action='store_true', help="Enable dynamic-range INT8 quantization")
    args = parser.parse_args()
    export_litert(args.model_dir, args.output, args.quantize)