| |
| """Exports moonshine.py's encoder and decoder to ExecuTorch 1.3.1 (the runtime anira's backends ship) and |
| checks them (check.py). |
| |
| uv run export_executorch.py <folder> <variant> [--coreml] <audio.wav>... |
| |
| writes <folder>/<variant>/executorch/{encoder,decoder}.pte, lowered to XNNPACK (CPU), float32; with |
| --coreml {encoder,decoder}_coreml.pte, lowered to Core ML (Apple platforms), float32. |
| """ |
|
|
| import os |
| import sys |
| import time |
|
|
| import coremltools as ct |
| import numpy as np |
| import torch |
| from executorch.backends.apple.coreml.compiler import CoreMLBackend |
| from executorch.backends.apple.coreml.partition import CoreMLPartitioner |
| from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner |
| from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower |
| from executorch.runtime import Runtime |
|
|
| import moonshine as m |
| from check import compare |
|
|
|
|
| def partitioner(coreml): |
| if not coreml: |
| return XnnpackPartitioner() |
| return CoreMLPartitioner(compile_specs=CoreMLBackend.generate_compile_specs( |
| compute_precision=ct.precision.FLOAT32, compute_unit=ct.ComputeUnit.ALL, |
| minimum_deployment_target=ct.target.macOS13)) |
|
|
|
|
| def main(): |
| folder, variant = sys.argv[1], sys.argv[2] |
| coreml = "--coreml" in sys.argv[3:] |
| files = [a for a in sys.argv[3:] if not a.startswith("--")] |
| model, tokenizer = m.load(variant) |
| os.makedirs(f"{folder}/{variant}/executorch", exist_ok=True) |
| modules = {"encoder": m.Encoder(model).eval(), "decoder": m.Decoder(model).eval()} |
| methods = {} |
| for graph, module in modules.items(): |
| start = time.perf_counter() |
| exported = torch.export.export(module, m.examples(model)[graph], strict=False) |
| edge = to_edge_transform_and_lower(exported, partitioner=[partitioner(coreml)], |
| compile_config=EdgeCompileConfig(_check_ir_validity=False)) |
| path = f"{folder}/{variant}/executorch/{graph}{'_coreml' if coreml else ''}.pte" |
| with open(path, "wb") as f: |
| f.write(edge.to_executorch().buffer) |
| print(f"{path}: {os.path.getsize(path):,} bytes, {time.perf_counter() - start:.0f} s", flush=True) |
| methods[graph] = Runtime.get().load_program(path).load_method("forward") |
|
|
| def run(graph, *values): |
| return [t.numpy() for t in methods[graph].execute([torch.from_numpy(np.ascontiguousarray(v)) for v in values])] |
|
|
| for graph in modules: |
| example = [x.numpy() for x in m.examples(model)[graph]] |
| run(graph, *example) |
| start = time.perf_counter() |
| run(graph, *example) |
| print(f"{graph}: {(time.perf_counter() - start) * 1e3:.2f} ms per call", flush=True) |
| compare("ExecuTorch Core ML" if coreml else "ExecuTorch", run, variant, files, model, tokenizer, |
| tolerance=1e-2 if coreml else 1e-3) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
