python/diag_size.py

import onnx, collections
from pathlib import Path
m = onnx.load("models/mvanet_box_segmenter.onnx")
inits = m.graph.initializer
def nbytes(t):
    import numpy as np
    return int(np.prod(t.dims)) * {1:4,2:1,3:1,4:2,5:2,6:4,7:8,10:2,11:8,9:1}.get(t.data_type,4)
tot = sum(nbytes(t) for t in inits)
print(f"initializers: {len(inits)}  total {tot/1e6:.1f} MB")
top = sorted(inits, key=nbytes, reverse=True)[:15]
for t in top:
    print(f"  {nbytes(t)/1e6:8.2f} MB  dtype={t.data_type}  dims={list(t.dims)}  {t.name[:70]}")
# count Constant nodes (node-embedded constants) too
cbytes=0; cn=0
for n in m.graph.node:
    if n.op_type=="Constant":
        cn+=1
        for a in n.attribute:
            if a.t.dims:
                import numpy as np
                cbytes += int(np.prod(a.t.dims))*{1:4,2:1,3:1,4:2,5:2,6:4,7:8,10:2,11:8,9:1}.get(a.t.data_type,4)
print(f"Constant nodes: {cn}  embedded ~{cbytes/1e6:.1f} MB")
# dtype histogram of initializers
h=collections.Counter(t.data_type for t in inits)
print("init dtype histogram (onnx enum):", dict(h))