File size: 2,232 Bytes
ce847d4 |
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 |
"""Quick inspection of specific ONNX models."""
import onnx
from pathlib import Path
for idx in [11, 12, 22, 33]:
matches = list(Path('oneocr_extracted/onnx_models').glob(f'model_{idx:02d}_*'))
if not matches:
continue
m = onnx.load(str(matches[0]))
print(f'\n=== model_{idx:02d} ({matches[0].name}) ===')
print('Opsets:', [(o.domain, o.version) for o in m.opset_import])
custom_ops = set()
for n in m.graph.node:
if n.domain:
custom_ops.add((n.domain, n.op_type))
print('Custom ops:', list(custom_ops))
for i in m.graph.input:
dims = []
if i.type.tensor_type.HasField('shape'):
for d in i.type.tensor_type.shape.dim:
v = str(d.dim_value) if d.dim_value else (d.dim_param or '?')
dims.append(v)
etype = i.type.tensor_type.elem_type
print(f' Input: {i.name} shape=[{", ".join(dims)}] dtype={etype}')
for o in m.graph.output:
dims = []
if o.type.tensor_type.HasField('shape'):
for d in o.type.tensor_type.shape.dim:
v = str(d.dim_value) if d.dim_value else (d.dim_param or '?')
dims.append(v)
etype = o.type.tensor_type.elem_type
print(f' Output: {o.name} shape=[{", ".join(dims)}] dtype={etype}')
cnt = {}
for n in m.graph.node:
key = f'{n.domain}::{n.op_type}' if n.domain else n.op_type
cnt[key] = cnt.get(key, 0) + 1
print(' Nodes:', dict(cnt))
# Show attributes for custom ops
for n in m.graph.node:
if n.domain:
attrs = {a.name: a for a in n.attribute}
attr_summary = {}
for k, a in attrs.items():
if a.type == 2: # FLOAT
attr_summary[k] = a.f
elif a.type == 1: # INT
attr_summary[k] = a.i
elif a.type == 3: # STRING
attr_summary[k] = a.s.decode()
print(f' CustomOp {n.op_type} attrs: {attr_summary}')
# Show input/output names
print(f' inputs: {list(n.input)}')
print(f' outputs: {list(n.output)}')
break # just show first custom op
|