File size: 3,109 Bytes
05b7816 | 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 | # Shrinks an exported BiRefNet: BiRefNet runs its backbone twice (full and half scale), and the export writes the backbone's
# weights once per call, some of them transposed. Identical initializers are merged, and a transposed copy becomes a
# Transpose of the original (ORT folds it back when the session is created, so it costs nothing per frame). Node names
# are made unique (the fp16 conversion repeats a few).
import sys, hashlib, onnx, numpy as np
from onnx import helper, numpy_helper
def digest(a):
return hashlib.sha1(np.ascontiguousarray(a).tobytes()).hexdigest()
def toposort(nodes, available):
"""The nodes in an order where each comes after what makes its inputs (the fp16 conversion appends its casts)."""
have, out, pending = set(available) | {''}, [], list(nodes)
while pending:
rest = []
for nd in pending:
if all(i in have for i in nd.input):
out.append(nd)
have.update(nd.output)
else:
rest.append(nd)
if len(rest) == len(pending): raise ValueError(f'{len(rest)} nodes have inputs nothing makes, e.g. {rest[0].name}')
pending = rest
return out
def shrink(m):
seen, rename, keep, transposed = {}, {}, [], []
for t in m.graph.initializer:
a = numpy_helper.to_array(t)
key = (t.data_type, a.shape, digest(a))
if key in seen:
rename[t.name] = seen[key]
continue
if a.ndim == 2 and a.size >= 4096:
kt = (t.data_type, a.T.shape, digest(a.T))
if kt in seen:
transposed.append(helper.make_node('Transpose', [seen[kt]], [t.name], perm=[1, 0], name=f'{t.name}_from_transpose'))
continue
seen[key] = t.name
keep.append(t)
del m.graph.initializer[:]
m.graph.initializer.extend(keep)
for nd in m.graph.node:
for i, n in enumerate(nd.input):
if n in rename: nd.input[i] = rename[n]
# The fp16 conversion casts a value once per op it keeps in fp32, under one name: identical casts collapse to the first.
made, unique = {}, []
for nd in list(m.graph.node):
prev = made.get(nd.output[0]) if nd.output else None
if prev is not None:
if prev.op_type != nd.op_type or list(prev.input) != list(nd.input): raise ValueError(f'two different nodes make {nd.output[0]}')
continue
for o in nd.output: made[o] = nd
unique.append(nd)
del m.graph.node[:]
m.graph.node.extend(unique)
nodes = toposort(transposed + list(m.graph.node), {i.name for i in m.graph.input} | {t.name for t in m.graph.initializer})
names = set()
for i, nd in enumerate(nodes):
if not nd.name or nd.name in names: nd.name = f'{nd.name or nd.op_type}_{i}'
names.add(nd.name)
del m.graph.node[:]
m.graph.node.extend(nodes)
return m, len(rename), len(transposed)
if __name__ == '__main__':
m, n, t = shrink(onnx.load(sys.argv[1]))
onnx.checker.check_model(m)
onnx.save(m, sys.argv[2])
print('merged', n, 'transposed', t)
|