| """Run a LiteRT graph from Python with the CompiledModel API (CPU, or GPU = Metal on the Mac), inputs and outputs by |
| signature name, shapes read from the flatbuffer.""" |
| import numpy as np |
|
|
|
|
| def sig_io(path): |
| """{signature key: {'inputs': [(name, shape)], 'outputs': [(name, shape)]}} from the flatbuffer.""" |
| from ai_edge_litert import schema_py_generated as schema |
| model = schema.ModelT.InitFromPackedBuf(open(path, 'rb').read(), 0) |
| out = {} |
| for sd in model.signatureDefs: |
| g = model.subgraphs[sd.subgraphIndex] |
| key = sd.signatureKey.decode() if isinstance(sd.signatureKey, bytes) else sd.signatureKey |
|
|
| def io(lst): |
| res = [] |
| for t in lst: |
| name = t.name.decode() if isinstance(t.name, bytes) else t.name |
| res.append((name, [int(s) for s in g.tensors[t.tensorIndex].shape])) |
| return res |
| out[key] = {'inputs': io(sd.inputs), 'outputs': io(sd.outputs)} |
| return out |
|
|
|
|
| class Graph: |
| def __init__(self, path, accel='gpu', enforce_f32=False): |
| from ai_edge_litert.compiled_model import CompiledModel |
| from ai_edge_litert.hardware_accelerator import HardwareAccelerator |
| from ai_edge_litert.options import GpuOptions, Options |
| if accel == 'gpu': |
| opts = Options(hardware_accelerators=HardwareAccelerator.GPU, gpu_options=GpuOptions(enforce_f32=enforce_f32)) |
| else: |
| opts = Options(hardware_accelerators=HardwareAccelerator.CPU) |
| self.path = path |
| self.m = CompiledModel.from_file(path, options=opts) |
| self.io = sig_io(path) |
| self.keys = list(self.io) |
| self.buffers = {} |
|
|
| def run(self, inputs, key=None): |
| """inputs: list in signature order or {name: array}. -> {output name: float32 array}.""" |
| key = key or self.keys[0] |
| idx = self.m.get_signature_index(key) |
| sig = self.m.get_signature_by_index(idx) |
| names_in = list(sig['inputs']) if isinstance(sig, dict) else [n for n, _ in self.io[key]['inputs']] |
| names_out = list(sig['outputs']) if isinstance(sig, dict) else [n for n, _ in self.io[key]['outputs']] |
| if key not in self.buffers: |
| self.buffers[key] = (self.m.create_input_buffers(idx), self.m.create_output_buffers(idx)) |
| ins, outs = self.buffers[key] |
| if not isinstance(inputs, dict): |
| inputs = dict(zip([n for n, _ in self.io[key]['inputs']], inputs)) |
| shapes_in = dict(self.io[key]['inputs']) |
| for name, buf in zip(names_in, ins): |
| arr = np.ascontiguousarray(np.asarray(inputs[name], dtype=np.float32)) |
| assert arr.size == int(np.prod(shapes_in[name])), (name, arr.shape, shapes_in[name]) |
| buf.write(arr.reshape(-1)) |
| self.m.run_by_index(idx, ins, outs) |
| shapes_out = dict(self.io[key]['outputs']) |
| res = {} |
| for j, (name, buf) in enumerate(zip(names_out, outs)): |
| n = int(np.prod(shapes_out[name])) |
| res[name] = buf.read(n, np.float32).reshape(shapes_out[name]).copy() |
| return res |
|
|
