| """Rewrite Gemm(transB=0, B=initializer) -> Gemm(transB=1, B=B^T) so ORT's CoreML EP can put the | |
| weight in weight.bin instead of inlining a transposed copy as hex text in model.mil. | |
| usage: fix_gemm.py in.onnx out.onnx""" | |
| import sys, onnx | |
| from onnx import numpy_helper | |
| m = onnx.load(sys.argv[1]) | |
| inits = {t.name: t for t in m.graph.initializer} | |
| n_fixed = 0 | |
| done = set() | |
| for n in m.graph.node: | |
| if n.op_type != "Gemm" or n.input[1] not in inits: | |
| continue | |
| attrs = {a.name: a for a in n.attribute} | |
| if "transB" in attrs and attrs["transB"].i == 1: | |
| continue | |
| name = n.input[1] + "_T" | |
| if name not in done: # weights shared by several Gemms (backbone runs twice) -> one copy | |
| w = numpy_helper.to_array(inits[n.input[1]]).T.copy() | |
| m.graph.initializer.append(numpy_helper.from_array(w, name)); done.add(name) | |
| n.input[1] = name | |
| if "transB" in attrs: | |
| attrs["transB"].i = 1 | |
| else: | |
| n.attribute.append(onnx.helper.make_attribute("transB", 1)) | |
| n_fixed += 1 | |
| used = {i for n in m.graph.node for i in n.input} | |
| keep = [t for t in m.graph.initializer if t.name in used] | |
| del m.graph.initializer[:] | |
| m.graph.initializer.extend(keep) | |
| onnx.save(m, sys.argv[2]) | |
| print("fixed", n_fixed, "Gemm nodes") | |