BiRefNet-lite matting, RTMW and YOLOX-tiny for onnxruntime-web on WebGPU

05b7816 verified 1 day ago

3.11 kB

# 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)

	# 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)