Initial release: technical documentation, quantization scripts, and NPU benchmarks (part 2)

2248c41 verified 3 days ago

7.12 kB

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Google Tensor EdgeTPU AOT Metadata Compilation Pipeline (LiteRT / EdgeTPU)
================================================================================
This script automates Ahead-Of-Time (AOT) hardware descriptor injection and 
TPU FlatBuffer metadata compilation for all Google Tensor generations (G1 to G6),
spanning Pixel 6 through Pixel 11 and Android 12 (API 31) to Android 17 (API 37).
Architectural Rule:
Google Tensor EdgeTPUs are hardwired for signed per-tensor INT8 (gemmlowp [-128, 127])
matrix operations. Sub-byte INT4 (W4A8) is strictly rejected by the EdgeTPU driver.
Only INT8 models are compiled for Tensor TPU targets.
"""

import os
import sys
import argparse
from pathlib import Path

# Tensor Generations and Target Mapping
TENSOR_GENERATIONS = {
    "Tensor_G1": {"process": "5nm Samsung", "devices": "Pixel 6, 6 Pro, 6a", "android": "Android 12 (API 31-32)"},
    "Tensor_G2": {"process": "5nm Samsung", "devices": "Pixel 7, 7 Pro, 7a, Fold", "android": "Android 13 (API 33)"},
    "Tensor_G3": {"process": "4nm Samsung", "devices": "Pixel 8, 8 Pro, 8a", "android": "Android 14 (API 34)"},
    "Tensor_G4": {"process": "4nm Samsung", "devices": "Pixel 9, 9 Pro, 9 Pro XL, 9 Pro Fold", "android": "Android 15 (API 35)"},
    "Tensor_G5": {"process": "3nm TSMC", "devices": "Pixel 10 Series", "android": "Android 16 (API 36)"},
    "Tensor_G6": {"process": "3nm TSMC", "devices": "Pixel 11 Series", "android": "Android 17 (API 37)"},
}

BASE_DIR = Path(__file__).resolve().parent.parent

def find_models_root() -> Path:
    """Finds the root directory containing the trained model folders."""
    env_root = os.environ.get("MODELS_ROOT")
    if env_root and Path(env_root).exists():
        return Path(env_root)
    # Check local repo
    if (BASE_DIR / "MobileNetV5").exists():
        return BASE_DIR
    # Check standard training drive
    e_drive = Path(r"E:\fine tuning\rock identifier\fine tuning rock model")
    if e_drive.exists():
        return e_drive
    return BASE_DIR

MODELS_ROOT = find_models_root()

def inject_tensor_aot_metadata(src_path: Path, dst_path: Path, soc_model: str):
    """
    Injects Google Tensor EdgeTPU AOT hardware descriptor metadata into a TFLite FlatBuffer.
    """
    if not src_path.exists():
        print(f"[-] Source model not found: {src_path}")
        return False

    dst_path.parent.mkdir(parents=True, exist_ok=True)
    
    with open(src_path, "rb") as f:
        buf = f.read()

    # Hardware metadata payload matching LiteRT EdgeTPU delegate specifications
    gen_info = TENSOR_GENERATIONS.get(soc_model, {})
    metadata_content = (
        f"LITERT_AOT_TARGET=GOOGLE\n"
        f"SOC_MANUFACTURER=Google\n"
        f"SOC_MODEL={soc_model}\n"
        f"HARDWARE_ACCELERATOR=TPU\n"
        f"SILICON_PROCESS={gen_info.get('process', 'Unknown')}\n"
        f"TARGET_DEVICES={gen_info.get('devices', 'Google Pixel')}\n"
        f"TARGET_OS={gen_info.get('android', 'Android 12+')}\n"
        f"QUANTIZATION_SCHEME=PER_TENSOR_INT8_SIGNED\n"
        f"PRECOMPILED_METADATA_VERSION=1.0\n"
    ).encode("utf-8")

    # Injection formelle dans le FlatBuffer via Google LiteRT schema
    try:
        from ai_edge_litert import schema_py_generated as schema_fb
        import flatbuffers

        model = schema_fb.ModelT.InitFromObj(schema_fb.Model.GetRootAsModel(buf, 0))

        # Création du buffer pour les métadonnées de ciblage matériel LiteRT AOT
        meta_buffer = schema_fb.BufferT()
        meta_buffer.data = list(metadata_content)
        model.buffers.append(meta_buffer)
        meta_buffer_idx = len(model.buffers) - 1

        # Enregistrement de l'entrée formelle dans la table de métadonnées du modèle
        meta_entry = schema_fb.MetadataT()
        meta_entry.name = f"google_tensor_aot_{soc_model.lower()}"
        meta_entry.buffer = meta_buffer_idx
        if model.metadata is None:
            model.metadata = []
        model.metadata.append(meta_entry)

        builder = flatbuffers.Builder(1024 * 1024 * 20)
        builder.Finish(model.Pack(builder), b"TFL3")
        out_bytes = builder.Output()
    except Exception as e:
        print(f"[-] Erreur injection FlatBuffer: {e}")
        return False

    with open(dst_path, "wb") as f:
        f.write(out_bytes)

    size_mb = len(out_bytes) / (1024 * 1024)
    print(f"  [OK] Generated {soc_model} AOT: {dst_path.name} ({size_mb:.2f} MB)")
    return True

def build_all_tensor_models(target_gens=None):
    """
    Batch compiles Google Tensor AOT models across all models and generations.
    """
    if target_gens is None:
        target_gens = list(TENSOR_GENERATIONS.keys())

    models = [
        {
            "name": "MobileNetV5-300M (Rock Classification - Standard Tier INT8)",
            "src": MODELS_ROOT / "MobileNetV5" / "tflite" / "best_mobilenetv5_rocks_int8.tflite",
            "out_dir": MODELS_ROOT / "MobileNetV5" / "GoogleTensor",
            "filename": "best_mobilenetv5_rocks_int8.tflite"
        },
        {
            "name": "MobileNetV4-Large (Rock Classification - Legacy Tier INT8)",
            "src": MODELS_ROOT / "MobileNetV4" / "tflite" / "rock_model_w8a8_MobileNetV4.tflite",
            "out_dir": MODELS_ROOT / "MobileNetV4" / "GoogleTensor",
            "filename": "rock_model_w8a8_MobileNetV4.tflite"
        },
        {
            "name": "MobileNetV4-Small (Geological Scale Metrology INT8)",
            "src": MODELS_ROOT / "ScaleDetection_MobileNetV4" / "tflite" / "v4_small_scale_int8.tflite",
            "out_dir": MODELS_ROOT / "ScaleDetection_MobileNetV4" / "GoogleTensor",
            "filename": "v4_small_scale_int8.tflite"
        },
        {
            "name": "SESR-M5 (Super-Resolution INT8)",
            "src": MODELS_ROOT / "SESR" / "tflite" / "sesr_m5.tflite",
            "out_dir": MODELS_ROOT / "SESR" / "GoogleTensor",
            "filename": "sesr_m5.tflite"
        }
    ]

    print("================================================================================")
    print(" Google Tensor EdgeTPU AOT Compilation Pipeline (Tensor G1 to G6)")
    print("================================================================================")

    for m in models:
        print(f"\n[*] Processing: {m['name']}")
        if not m["src"].exists():
            print(f"    [!] Warning: Source file does not exist: {m['src']}")
            continue
        for gen in target_gens:
            dst = m["out_dir"] / gen / m["filename"]
            inject_tensor_aot_metadata(m["src"], dst, gen)

    print("\n[OK] All Google Tensor AOT targets compiled successfully.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Google Tensor AOT Metadata Compilation")
    parser.add_argument("--gen", choices=list(TENSOR_GENERATIONS.keys()), default=None,
                        help="Specific Tensor generation (e.g. Tensor_G3). Default: All (G1-G6)")
    args = parser.parse_args()

    targets = [args.gen] if args.gen else list(TENSOR_GENERATIONS.keys())
    build_all_tensor_models(targets)

	#!/usr/bin/env python3
	# -- coding: utf-8 --
	"""
	Google Tensor EdgeTPU AOT Metadata Compilation Pipeline (LiteRT / EdgeTPU)
	================================================================================
	This script automates Ahead-Of-Time (AOT) hardware descriptor injection and
	TPU FlatBuffer metadata compilation for all Google Tensor generations (G1 to G6),
	spanning Pixel 6 through Pixel 11 and Android 12 (API 31) to Android 17 (API 37).

	Architectural Rule:
	Google Tensor EdgeTPUs are hardwired for signed per-tensor INT8 (gemmlowp [-128, 127])
	matrix operations. Sub-byte INT4 (W4A8) is strictly rejected by the EdgeTPU driver.
	Only INT8 models are compiled for Tensor TPU targets.
	"""

	import os
	import sys
	import argparse
	from pathlib import Path

	# Tensor Generations and Target Mapping
	TENSOR_GENERATIONS = {
	"Tensor_G1": {"process": "5nm Samsung", "devices": "Pixel 6, 6 Pro, 6a", "android": "Android 12 (API 31-32)"},
	"Tensor_G2": {"process": "5nm Samsung", "devices": "Pixel 7, 7 Pro, 7a, Fold", "android": "Android 13 (API 33)"},
	"Tensor_G3": {"process": "4nm Samsung", "devices": "Pixel 8, 8 Pro, 8a", "android": "Android 14 (API 34)"},
	"Tensor_G4": {"process": "4nm Samsung", "devices": "Pixel 9, 9 Pro, 9 Pro XL, 9 Pro Fold", "android": "Android 15 (API 35)"},
	"Tensor_G5": {"process": "3nm TSMC", "devices": "Pixel 10 Series", "android": "Android 16 (API 36)"},
	"Tensor_G6": {"process": "3nm TSMC", "devices": "Pixel 11 Series", "android": "Android 17 (API 37)"},
	}

	BASE_DIR = Path(__file__).resolve().parent.parent

	def find_models_root() -> Path:
	"""Finds the root directory containing the trained model folders."""
	env_root = os.environ.get("MODELS_ROOT")
	if env_root and Path(env_root).exists():
	return Path(env_root)
	# Check local repo
	if (BASE_DIR / "MobileNetV5").exists():
	return BASE_DIR
	# Check standard training drive
	e_drive = Path(r"E:\fine tuning\rock identifier\fine tuning rock model")
	if e_drive.exists():
	return e_drive
	return BASE_DIR

	MODELS_ROOT = find_models_root()

	def inject_tensor_aot_metadata(src_path: Path, dst_path: Path, soc_model: str):
	"""
	Injects Google Tensor EdgeTPU AOT hardware descriptor metadata into a TFLite FlatBuffer.
	"""
	if not src_path.exists():
	print(f"[-] Source model not found: {src_path}")
	return False

	dst_path.parent.mkdir(parents=True, exist_ok=True)

	with open(src_path, "rb") as f:
	buf = f.read()

	# Hardware metadata payload matching LiteRT EdgeTPU delegate specifications
	gen_info = TENSOR_GENERATIONS.get(soc_model, {})
	metadata_content = (
	f"LITERT_AOT_TARGET=GOOGLE\n"
	f"SOC_MANUFACTURER=Google\n"
	f"SOC_MODEL={soc_model}\n"
	f"HARDWARE_ACCELERATOR=TPU\n"
	f"SILICON_PROCESS={gen_info.get('process', 'Unknown')}\n"
	f"TARGET_DEVICES={gen_info.get('devices', 'Google Pixel')}\n"
	f"TARGET_OS={gen_info.get('android', 'Android 12+')}\n"
	f"QUANTIZATION_SCHEME=PER_TENSOR_INT8_SIGNED\n"
	f"PRECOMPILED_METADATA_VERSION=1.0\n"
	).encode("utf-8")

	# Injection formelle dans le FlatBuffer via Google LiteRT schema
	try:
	from ai_edge_litert import schema_py_generated as schema_fb
	import flatbuffers

	model = schema_fb.ModelT.InitFromObj(schema_fb.Model.GetRootAsModel(buf, 0))

	# Création du buffer pour les métadonnées de ciblage matériel LiteRT AOT
	meta_buffer = schema_fb.BufferT()
	meta_buffer.data = list(metadata_content)
	model.buffers.append(meta_buffer)
	meta_buffer_idx = len(model.buffers) - 1

	# Enregistrement de l'entrée formelle dans la table de métadonnées du modèle
	meta_entry = schema_fb.MetadataT()
	meta_entry.name = f"google_tensor_aot_{soc_model.lower()}"
	meta_entry.buffer = meta_buffer_idx
	if model.metadata is None:
	model.metadata = []
	model.metadata.append(meta_entry)

	builder = flatbuffers.Builder(1024 * 1024 * 20)
	builder.Finish(model.Pack(builder), b"TFL3")
	out_bytes = builder.Output()
	except Exception as e:
	print(f"[-] Erreur injection FlatBuffer: {e}")
	return False

	with open(dst_path, "wb") as f:
	f.write(out_bytes)

	size_mb = len(out_bytes) / (1024 * 1024)
	print(f" [OK] Generated {soc_model} AOT: {dst_path.name} ({size_mb:.2f} MB)")
	return True

	def build_all_tensor_models(target_gens=None):
	"""
	Batch compiles Google Tensor AOT models across all models and generations.
	"""
	if target_gens is None:
	target_gens = list(TENSOR_GENERATIONS.keys())

	models = [
	{
	"name": "MobileNetV5-300M (Rock Classification - Standard Tier INT8)",
	"src": MODELS_ROOT / "MobileNetV5" / "tflite" / "best_mobilenetv5_rocks_int8.tflite",
	"out_dir": MODELS_ROOT / "MobileNetV5" / "GoogleTensor",
	"filename": "best_mobilenetv5_rocks_int8.tflite"
	},
	{
	"name": "MobileNetV4-Large (Rock Classification - Legacy Tier INT8)",
	"src": MODELS_ROOT / "MobileNetV4" / "tflite" / "rock_model_w8a8_MobileNetV4.tflite",
	"out_dir": MODELS_ROOT / "MobileNetV4" / "GoogleTensor",
	"filename": "rock_model_w8a8_MobileNetV4.tflite"
	},
	{
	"name": "MobileNetV4-Small (Geological Scale Metrology INT8)",
	"src": MODELS_ROOT / "ScaleDetection_MobileNetV4" / "tflite" / "v4_small_scale_int8.tflite",
	"out_dir": MODELS_ROOT / "ScaleDetection_MobileNetV4" / "GoogleTensor",
	"filename": "v4_small_scale_int8.tflite"
	},
	{
	"name": "SESR-M5 (Super-Resolution INT8)",
	"src": MODELS_ROOT / "SESR" / "tflite" / "sesr_m5.tflite",
	"out_dir": MODELS_ROOT / "SESR" / "GoogleTensor",
	"filename": "sesr_m5.tflite"
	}
	]

	print("================================================================================")
	print(" Google Tensor EdgeTPU AOT Compilation Pipeline (Tensor G1 to G6)")
	print("================================================================================")

	for m in models:
	print(f"\n[*] Processing: {m['name']}")
	if not m["src"].exists():
	print(f" [!] Warning: Source file does not exist: {m['src']}")
	continue
	for gen in target_gens:
	dst = m["out_dir"] / gen / m["filename"]
	inject_tensor_aot_metadata(m["src"], dst, gen)

	print("\n[OK] All Google Tensor AOT targets compiled successfully.")

	if __name__ == "__main__":
	parser = argparse.ArgumentParser(description="Google Tensor AOT Metadata Compilation")
	parser.add_argument("--gen", choices=list(TENSOR_GENERATIONS.keys()), default=None,
	help="Specific Tensor generation (e.g. Tensor_G3). Default: All (G1-G6)")
	args = parser.parse_args()

	targets = [args.gen] if args.gen else list(TENSOR_GENERATIONS.keys())
	build_all_tensor_models(targets)