File size: 3,285 Bytes
fdefdba | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
"""
ONNX graph surgery to fix det_500m.onnx for batch-independent inference.
Root cause: 9 output Transpose nodes use perm [2,3,0,1], moving [N,C,H,W] to
[H,W,N,C]. The subsequent Reshape [-1,K] flattens everything, interleaving frames.
Fix:
1. Transpose perm [2,3,0,1] -> [0,2,3,1]: keeps batch first [N,H,W,C]
2. Reshape [-1,K] -> [0,-1,K]: preserves batch dim -> [N, anchors, K]
3. Input/output declarations updated for dynamic batch
"""
import argparse
import onnx
from onnx import numpy_helper, TensorProto, helper
import numpy as np
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True)
parser.add_argument('--out', required=True)
args = parser.parse_args()
model = onnx.load(args.model)
# 1. Make input batch dim dynamic
for inp in model.graph.input:
d0 = inp.type.tensor_type.shape.dim[0]
d0.ClearField('dim_value')
d0.dim_param = 'batch'
print(f"Input '{inp.name}': batch dim -> dynamic")
# 2. Fix the 9 output Transpose nodes: perm [2,3,0,1] -> [0,2,3,1]
output_transpose_names = {
'Transpose_154', 'Transpose_158', 'Transpose_161', # stride 8
'Transpose_179', 'Transpose_183', 'Transpose_186', # stride 16
'Transpose_204', 'Transpose_208', 'Transpose_211', # stride 32
}
patched_transposes = 0
for node in model.graph.node:
if node.name in output_transpose_names:
for attr in node.attribute:
if attr.name == 'perm':
old_perm = list(attr.ints)
attr.ints[:] = [0, 2, 3, 1]
print(f" {node.name}: perm {old_perm} -> [0, 2, 3, 1]")
patched_transposes += 1
print(f"Patched {patched_transposes} Transpose nodes")
# 3. Fix Reshape initializers: [-1, K] -> [0, -1, K]
# These are the output reshape shapes that flatten spatial+batch together
output_reshape_inits = {'441', '445', '448'}
patched_reshapes = 0
for init in list(model.graph.initializer):
if init.name not in output_reshape_inits:
continue
arr = numpy_helper.to_array(init).copy()
print(f" Reshape init '{init.name}': {arr} -> ", end='')
new_arr = np.array([0] + arr.tolist(), dtype=np.int64)
print(new_arr)
new_t = numpy_helper.from_array(new_arr, init.name)
model.graph.initializer.remove(init)
model.graph.initializer.append(new_t)
patched_reshapes += 1
print(f"Patched {patched_reshapes} Reshape initializer(s)")
# 4. Fix output shapes: 2D [anchors, K] -> 3D [batch, anchors, K]
for out in model.graph.output:
shape = out.type.tensor_type.shape
old_dims = []
for d in shape.dim:
if d.HasField('dim_param'):
old_dims.append(('param', d.dim_param))
else:
old_dims.append(('value', d.dim_value))
while len(shape.dim) > 0:
shape.dim.pop()
# Add batch dim
d = shape.dim.add()
d.dim_param = 'batch'
# Re-add original dims
for kind, val in old_dims:
d = shape.dim.add()
if kind == 'param':
d.dim_param = val
else:
d.dim_value = val
new_shape = [dd.dim_param if dd.HasField('dim_param') else dd.dim_value for dd in shape.dim]
print(f"Output '{out.name}': {old_dims} -> {new_shape}")
# Save
onnx.save(model, args.out)
print(f"\nSaved: {args.out}")
|