v0.7.0: 20-60% less memory (own blob store, vision loaded on demand and freed, host tables mapped, direct-I/O expert load); REAP streams too; first-on-Intel-NPU SSD streaming note
4955f2e verified | """Does importing blobs ourselves (core.import_model, no CACHE_DIR) avoid keeping the blob next to the driver's copy? | |
| python mem_import.py export MODEL_DIR S OUT_DIR compile (cache hit) and export every S-segment's blob | |
| python mem_import.py import OUT_DIR MODE import them and measure; MODE = bytes | aligned | unaligned""" | |
| import io, json, mmap, os, sys, time | |
| import numpy as np | |
| import openvino as ov | |
| import psutil | |
| P = psutil.Process() | |
| def mem(): | |
| m = P.memory_info() | |
| return m.rss / 2**30, getattr(m, "private", 0) / 2**30 | |
| def export(d, S, out): | |
| from onw.runtime import read_segment | |
| meta = json.load(open(os.path.join(d, "engine.json"))) | |
| core = ov.Core() | |
| os.makedirs(out, exist_ok=True) | |
| for n, m in sorted(meta["segments"].items()): | |
| if m["S"] != S: | |
| continue | |
| c = core.compile_model(read_segment(core, d, n, m, False, {}, meta["T"], meta["T"]), "NPU", | |
| {"CACHE_DIR": os.path.join(d, "npu_cache")}) | |
| s = io.BytesIO() | |
| c.export_model(s) | |
| open(os.path.join(out, n[:-4] + ".blob"), "wb").write(s.getvalue()) | |
| print("exported", len(os.listdir(out)), "blobs,", sum(os.path.getsize(os.path.join(out, f)) for f in os.listdir(out)) / 2**30, "GB") | |
| def imp(out, mode): | |
| core = ov.Core() | |
| marks = [("start", *mem())] | |
| compiled = [] | |
| for f in sorted(os.listdir(out)): | |
| path = os.path.join(out, f) | |
| n = os.path.getsize(path) | |
| if mode == "bytes": | |
| data = open(path, "rb").read() | |
| c = core.import_model(io.BytesIO(data), "NPU", {}) | |
| del data | |
| else: # an ov.Tensor over our own buffer (page aligned or not) | |
| buf = mmap.mmap(-1, n + 4096) | |
| off = 0 if mode == "aligned" else 64 | |
| with open(path, "rb") as fh: | |
| fh.readinto(memoryview(buf)[off:off + n]) | |
| t = ov.Tensor(np.frombuffer(buf, np.uint8, n, off), shared_memory=True) | |
| c = core.import_model(t, "NPU", {}) | |
| del t | |
| buf.close() | |
| compiled.append(c) | |
| marks.append((f"import_model ({mode})", *mem())) | |
| import gc; gc.collect() | |
| marks.append(("gc", *mem())) | |
| reqs = [c.create_infer_request() for c in compiled] | |
| ctx, shared = core.get_default_context("NPU"), {} | |
| for r in reqs: # every input bound by us, as onw does | |
| for i in r.model_inputs: | |
| shp = [x.get_length() for x in i.get_partial_shape()] | |
| key = (str(i.get_element_type()), tuple(shp)) | |
| if key not in shared: | |
| shared[key] = ctx.create_host_tensor(i.get_element_type(), ov.Shape(shp)) | |
| r.set_tensor(i.get_any_name(), shared[key]) | |
| for r in reqs: | |
| r.infer() | |
| marks.append(("create_infer_request + first infer", *mem())) | |
| prev = 0 | |
| for name, w, p in marks: | |
| print(f" {name:40s} working set {w:6.2f} GB ({w - prev:+5.2f}) private {p:5.2f} GB") | |
| prev = w | |
| if __name__ == "__main__": | |
| if sys.argv[1] == "export": | |
| export(sys.argv[2], int(sys.argv[3]), sys.argv[4]) | |
| else: | |
| imp(sys.argv[2], sys.argv[3]) | |