File size: 7,121 Bytes
2248c41 | 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 164 165 166 167 168 | #!/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)
|