File size: 3,147 Bytes
ed66dc1 | 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 | """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])) # the Python API takes (output_index, signature_index)
res[name] = buf.read(n, np.float32).reshape(shapes_out[name]).copy()
return res
|