File size: 4,698 Bytes
9496f98 | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | #!/usr/bin/env python3
"""Inspect an ONNX model and emit a machine-readable characterization.
Reports the facts a production-readiness review needs about a checkpoint before and
after optimization: opset, node/initializer counts, the weight-dtype histogram, whether
external data is used, the exact I/O tensor contract (names, dtypes, shapes), and the
operator histogram. Pure `onnx` (no onnxruntime) so it also runs on graphs a given ORT
build cannot execute.
Usage:
python inspect_onnx.py MODEL.onnx [--json OUT.json]
"""
from __future__ import annotations
import argparse
import collections
import hashlib
import json
import os
import sys
import onnx
from onnx import TensorProto
# TensorProto enum value -> human dtype name, for the weight/IO histograms.
DTYPE_NAME = {v: k for k, v in TensorProto.DataType.items()}
def _elem_type_name(elem_type: int) -> str:
return DTYPE_NAME.get(elem_type, f"UNKNOWN({elem_type})")
def _shape_of(value_info) -> list:
dims = []
for d in value_info.type.tensor_type.shape.dim:
if d.HasField("dim_value"):
dims.append(d.dim_value)
elif d.HasField("dim_param") and d.dim_param:
dims.append(d.dim_param)
else:
dims.append("dynamic")
return dims
def _io_spec(value_info) -> dict:
return {
"name": value_info.name,
"dtype": _elem_type_name(value_info.type.tensor_type.elem_type),
"shape": _shape_of(value_info),
}
def inspect(path: str) -> dict:
size_bytes = os.path.getsize(path)
with open(path, "rb") as fh:
raw = fh.read()
sha256 = hashlib.sha256(raw).hexdigest()
model = onnx.load(path, load_external_data=False)
graph = model.graph
# Opset per domain.
opset = {}
for entry in model.opset_import:
opset[entry.domain or "ai.onnx"] = entry.version
# Operator histogram.
op_hist = collections.Counter(node.op_type for node in graph.node)
# Initializer dtype histogram + external-data detection.
init_dtype_hist = collections.Counter()
init_bytes = collections.Counter()
uses_external_data = False
for init in graph.initializer:
dtype = _elem_type_name(init.data_type)
init_dtype_hist[dtype] += 1
# A tensor stored externally has data_location == EXTERNAL(1).
if init.data_location == TensorProto.EXTERNAL:
uses_external_data = True
# Rough in-graph byte estimate (raw_data if present).
init_bytes[dtype] += len(init.raw_data)
# Sidecar external-data files sitting next to the model.
external_files = []
model_dir = os.path.dirname(os.path.abspath(path))
for entry in os.listdir(model_dir) if os.path.isdir(model_dir) else []:
if entry.endswith(".onnx_data") or entry.endswith(".data") or entry.endswith(".bin"):
external_files.append(entry)
# I/O contract, skipping initializers that also appear as graph inputs.
init_names = {init.name for init in graph.initializer}
inputs = [_io_spec(v) for v in graph.input if v.name not in init_names]
outputs = [_io_spec(v) for v in graph.output]
checker_ok = True
checker_error = None
try:
onnx.checker.check_model(model)
except Exception as exc: # noqa: BLE001 - report any checker failure verbatim
checker_ok = False
checker_error = str(exc)
return {
"path": os.path.abspath(path),
"size_bytes": size_bytes,
"size_mib": round(size_bytes / 2**20, 3),
"sha256": sha256,
"ir_version": model.ir_version,
"producer": f"{model.producer_name} {model.producer_version}".strip(),
"opset": opset,
"num_nodes": len(graph.node),
"num_initializers": len(graph.initializer),
"uses_external_data": uses_external_data,
"external_files_in_dir": external_files,
"initializer_dtype_histogram": dict(init_dtype_hist),
"initializer_bytes_by_dtype": dict(init_bytes),
"operator_histogram": dict(sorted(op_hist.items(), key=lambda kv: -kv[1])),
"num_distinct_ops": len(op_hist),
"inputs": inputs,
"outputs": outputs,
"onnx_checker_ok": checker_ok,
"onnx_checker_error": checker_error,
}
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("model")
ap.add_argument("--json", dest="json_out", default=None)
args = ap.parse_args()
info = inspect(args.model)
text = json.dumps(info, indent=2)
print(text)
if args.json_out:
with open(args.json_out, "w") as fh:
fh.write(text)
return 0
if __name__ == "__main__":
sys.exit(main())
|