File size: 2,721 Bytes
d90d22d | 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 | # SPDX-License-Identifier: Apache-2.0
"""Exports moonshine.py's encoder and decoder to Core ML ML programs (.mlpackage, macOS 13 / iOS 16,
float32) with the ONNX export's tensor names, and checks them (check.py).
uv run export_coreml.py <folder> <variant> <audio.wav>...
writes <folder>/<variant>/coreml/{encoder,decoder}.mlpackage. Token, positions and lengths are int32.
"""
import sys
import time
import coremltools as ct
import numpy as np
import torch
import moonshine as m
from check import compare
def drop_aliases(exported):
"""aten.alias is a no-op that coremltools cannot convert."""
graph = exported.graph_module.graph
for node in list(graph.nodes):
if node.op == "call_function" and node.target == torch.ops.aten.alias.default:
node.replace_all_uses_with(node.args[0])
graph.erase_node(node)
exported.graph_module.recompile()
def main():
folder, variant, files = sys.argv[1], sys.argv[2], sys.argv[3:]
model, tokenizer = m.load(variant)
modules = {"encoder": m.Encoder(model).eval(), "decoder": m.Decoder(model).eval()}
converted = {}
for graph, module in modules.items():
start = time.perf_counter()
example = m.examples(model)[graph]
exported = torch.export.export(module, example, strict=False).run_decompositions({})
drop_aliases(exported)
inputs, outputs = m.NAMES[graph]
converted[graph] = ct.convert(
exported, convert_to="mlprogram", minimum_deployment_target=ct.target.macOS13,
compute_precision=ct.precision.FLOAT32,
inputs=[ct.TensorType(name=n, shape=x.shape, dtype=np.float32 if x.is_floating_point() else np.int32)
for n, x in zip(inputs, example)],
outputs=[ct.TensorType(name=n) for n in outputs])
converted[graph].save(f"{folder}/{variant}/coreml/{graph}.mlpackage")
print(f"{folder}/{variant}/coreml/{graph}.mlpackage: {time.perf_counter() - start:.0f} s", flush=True)
def run(graph, *values):
inputs, outputs = m.NAMES[graph]
feed = {n: np.asarray(v).astype(np.int32 if np.asarray(v).dtype.kind == "i" else np.float32)
for n, v in zip(inputs, values)}
result = converted[graph].predict(feed)
return [result[n] for n in outputs]
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("Core ML", run, variant, files, model, tokenizer, tolerance=1e-2)
if __name__ == "__main__":
main()
|