#!/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)