File size: 7,516 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 169 170 171 172 173 174 175 176 177 178 179 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MediaTek NeuroPilot APU AOT Metadata Compilation Pipeline (LiteRT)
================================================================================
Ce script automatise la compilation et l'injection de métadonnées AOT LiteRT
pour les processeurs MediaTek Dimensity et leurs APU dédiées.
Règles matérielles :
- Dimensity 8300 / 9300 / 9400 (APU 780, 790, 890) -> W4A8 prioritaire (+ INT8 standard).
- Dimensity 9000 / 9200 (APU 5.0, 680) -> Strictement INT8 pur (pas de bloc sub-byte matériel).
- SESR-M5 Super-Resolution -> Strictement INT8 sur toutes les cibles.
"""
import os
import sys
import argparse
from pathlib import Path
MEDIATEK_GENERATIONS = {
"MT6983": {
"name": "Dimensity 9000 / 9000+",
"apu": "APU 5.0",
"precision": "INT8",
"devices": "OnePlus Nord 3, Vivo X80",
"android": "Android 12-14"
},
"MT6985": {
"name": "Dimensity 9200 / 9200+",
"apu": "APU 680",
"precision": "INT8",
"devices": "Vivo X90, Xiaomi 13T Pro",
"android": "Android 13-15"
},
"MT6897": {
"name": "Dimensity 8300 / 8300-Ultra",
"apu": "APU 780",
"precision": "W4A8",
"devices": "POCO X6 Pro",
"android": "Android 14-16"
},
"MT6989": {
"name": "Dimensity 9300 / 9300+",
"apu": "APU 790",
"precision": "W4A8",
"devices": "Vivo X100 Pro, Pad Pro",
"android": "Android 14-16"
},
"MT6991": {
"name": "Dimensity 9400",
"apu": "APU 890",
"precision": "W4A8",
"devices": "Vivo X200 Series, Oppo Find X8",
"android": "Android 15-17"
}
}
BASE_DIR = Path(__file__).resolve().parent
def find_models_root() -> Path:
env_root = os.environ.get("MODELS_ROOT")
if env_root and Path(env_root).exists():
return Path(env_root)
if (BASE_DIR / "MobileNetV5").exists():
return BASE_DIR
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_mediatek_metadata(src_path: Path, dst_path: Path, soc_model: str, gen_info: dict, quant_type: str):
if not src_path.exists():
print(f" [-] Source model missing: {src_path}")
return False
dst_path.parent.mkdir(parents=True, exist_ok=True)
with open(src_path, "rb") as f:
buf = f.read()
metadata_content = (
f"LITERT_AOT_TARGET=MEDIATEK\n"
f"SOC_MANUFACTURER=MediaTek\n"
f"SOC_MODEL={soc_model}\n"
f"SOC_MARKETING_NAME={gen_info.get('name', '')}\n"
f"HARDWARE_ACCELERATOR={gen_info.get('apu', 'APU')}\n"
f"TARGET_DEVICES={gen_info.get('devices', '')}\n"
f"TARGET_OS={gen_info.get('android', 'Android 14+')}\n"
f"QUANTIZATION_SCHEME={quant_type}\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))
meta_buffer = schema_fb.BufferT()
meta_buffer.data = list(metadata_content)
model.buffers.append(meta_buffer)
meta_buffer_idx = len(model.buffers) - 1
meta_entry = schema_fb.MetadataT()
meta_entry.name = f"mediatek_neuropilot_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 MediaTek: {e}")
return False
with open(dst_path, "wb") as f:
f.write(out_bytes)
size_mb = len(out_bytes) / (1024 * 1024)
print(f" [OK] MediaTek [{soc_model}] {dst_path.name} ({size_mb:.2f} MB)")
return True
def compile_mediatek_targets(targets=None):
if targets is None:
targets = list(MEDIATEK_GENERATIONS.keys())
print("=" * 80)
print(" Compiling MediaTek NeuroPilot AOT Targets (Dimensity)")
print("=" * 80)
for soc_model in targets:
if soc_model not in MEDIATEK_GENERATIONS:
print(f"[!] Target inconnue: {soc_model}")
continue
info = MEDIATEK_GENERATIONS[soc_model]
is_w4 = (info["precision"] == "W4A8")
print(f"\n[*] Target: MediaTek {soc_model} ({info['name']} - {info['apu']}) -> Priorite: {info['precision']}")
# 1. MobileNetV5
v5_src = MODELS_ROOT / "MobileNetV5" / "tflite" / ("best_mobilenetv5_rocks_w4a8.tflite" if is_w4 else "best_mobilenetv5_rocks_int8.tflite")
v5_out = MODELS_ROOT / "MobileNetV5" / "MediaTek" / soc_model / ("best_mobilenetv5_rocks_w4a8.tflite" if is_w4 else "best_mobilenetv5_rocks_int8.tflite")
inject_mediatek_metadata(v5_src, v5_out, soc_model, info, "W4A8_FULL_INTEGER" if is_w4 else "PER_TENSOR_INT8_SIGNED")
if is_w4:
v5_int8_src = MODELS_ROOT / "MobileNetV5" / "tflite" / "best_mobilenetv5_rocks_int8.tflite"
v5_int8_out = MODELS_ROOT / "MobileNetV5" / "MediaTek" / soc_model / "best_mobilenetv5_rocks_int8.tflite"
inject_mediatek_metadata(v5_int8_src, v5_int8_out, soc_model, info, "PER_TENSOR_INT8_SIGNED")
# 2. MobileNetV4
v4_src = MODELS_ROOT / "MobileNetV4" / "tflite" / ("rock_model_w4a8_MobileNetV4.tflite" if is_w4 else "rock_model_w8a8_MobileNetV4.tflite")
v4_out = MODELS_ROOT / "MobileNetV4" / "MediaTek" / soc_model / ("rock_model_w4a8_MobileNetV4.tflite" if is_w4 else "rock_model_w8a8_MobileNetV4.tflite")
inject_mediatek_metadata(v4_src, v4_out, soc_model, info, "W4A8_FULL_INTEGER" if is_w4 else "PER_TENSOR_INT8_SIGNED")
# 3. Scale Detection
scale_src = MODELS_ROOT / "ScaleDetection_MobileNetV4" / "tflite" / ("v4_small_scale_int4.tflite" if is_w4 else "v4_small_scale_int8.tflite")
scale_out = MODELS_ROOT / "ScaleDetection_MobileNetV4" / "MediaTek" / soc_model / ("v4_small_scale_int4.tflite" if is_w4 else "v4_small_scale_int8.tflite")
inject_mediatek_metadata(scale_src, scale_out, soc_model, info, "W4A8_FULL_INTEGER" if is_w4 else "PER_TENSOR_INT8_SIGNED")
# 4. SESR Super-Resolution (Strictement INT8)
sesr_src = MODELS_ROOT / "SESR" / "tflite" / "sesr_m5.tflite"
sesr_out = MODELS_ROOT / "SESR" / "MediaTek" / soc_model / "sesr_m5.tflite"
inject_mediatek_metadata(sesr_src, sesr_out, soc_model, info, "PER_TENSOR_INT8_SIGNED")
print("\n[OK] Compilation MediaTek terminee avec succes.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MediaTek NeuroPilot AOT Compilation")
parser.add_argument("--soc", choices=list(MEDIATEK_GENERATIONS.keys()), default=None,
help="Cible SoC specifique (ex: MT6989). Par defaut: Toutes les cibles.")
args = parser.parse_args()
targets = [args.soc] if args.soc else list(MEDIATEK_GENERATIONS.keys())
compile_mediatek_targets(targets)
|