File size: 7,705 Bytes
001425d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | #!/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)
|