File size: 3,703 Bytes
3f1985b | 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 | #!/usr/bin/env python3
"""
det02 — ONNX graph surgery to patch det_500m.onnx for dynamic batch.
Patches 3 things in the graph metadata:
1. Input dim[0]: 1 → symbolic 'batch'
2. Reshape shape constants starting with 1: that 1 → 0 (copy batch dim from input)
3. Output dim[0]: 1 → symbolic 'batch'
Saves the result, then prints declared output shapes and actual runtime shapes for N=1 and N=2.
"""
import argparse
import numpy as np
import onnx
from onnx import numpy_helper
import onnxruntime as ort
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to det_500m.onnx')
parser.add_argument('--out', required=True, help='Output path for patched model')
args = parser.parse_args()
model = onnx.load(args.model)
# ── 1. Patch input batch dim ──────────────────────────────────────────────────
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. Patch Reshape shape initializers that start with 1 ────────────────────
reshape_shape_names = set()
for node in model.graph.node:
if node.op_type == 'Reshape':
reshape_shape_names.add(node.input[1])
patched = 0
for init in list(model.graph.initializer):
if init.name not in reshape_shape_names:
continue
arr = numpy_helper.to_array(init).copy()
if arr.ndim == 1 and arr[0] == 1:
print(f" Reshape initializer '{init.name}': {arr} → ", end='')
arr[0] = 0
print(arr)
new_t = numpy_helper.from_array(arr, init.name)
model.graph.initializer.remove(init)
model.graph.initializer.append(new_t)
patched += 1
print(f' Patched {patched} Reshape shape constant(s)')
# ── 3. Patch output batch dims ────────────────────────────────────────────────
for out in model.graph.output:
d0 = out.type.tensor_type.shape.dim[0]
d0.ClearField('dim_value')
d0.dim_param = 'batch'
print(f" Output '{out.name}': batch dim → dynamic")
# ── Save ──────────────────────────────────────────────────────────────────────
onnx.checker.check_model(model)
onnx.save(model, args.out)
print(f'\nSaved: {args.out}')
# ── Declared output shapes (from patched graph) ───────────────────────────────
print('\nDeclared output shapes after surgery:')
for out in model.graph.output:
shape = [d.dim_param if d.HasField('dim_param') else d.dim_value
for d in out.type.tensor_type.shape.dim]
print(f' {out.name}: {shape}')
# ── Actual output shapes (from runtime with N=1 and N=2) ─────────────────────
print('\nActual output shapes at runtime:')
session = ort.InferenceSession(args.out, providers=['CPUExecutionProvider'])
inp0 = session.get_inputs()[0]
out_names = [o.name for o in session.get_outputs()]
for N in (1, 2):
dummy = np.random.randint(0, 255, (N, 3, 640, 640)).astype(np.float32)
try:
outs = session.run(out_names, {inp0.name: dummy})
print(f' N={N}:')
for name, o in zip(out_names, outs):
print(f' {name}: {list(o.shape)}')
except Exception as e:
print(f' N={N}: FAILED — {e}')
|