File size: 2,456 Bytes
fbcc547 | 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 | import onnx, sys
from onnx import helper
src, dst, mode = sys.argv[1], sys.argv[2], sys.argv[3] # mode: pads | static:H,W
m = onnx.load(src)
# 1) auto_pad SAME_UPPER -> explicit pads (kernel k, stride 1 => total pad k-1, upper side)
n_fixed = 0
for n in m.graph.node:
if n.op_type in ("MaxPool", "AveragePool", "Conv"):
ap = [a for a in n.attribute if a.name == "auto_pad"]
if ap and ap[0].s in (b"SAME_UPPER", b"SAME_LOWER"):
ks = [a for a in n.attribute if a.name == "kernel_shape"][0].ints
st = [a for a in n.attribute if a.name == "strides"]
st = st[0].ints if st else [1]*len(ks)
assert all(s == 1 for s in st), (n.name, list(st))
tot = [k-1 for k in ks]
if ap[0].s == b"SAME_UPPER":
pads = [t//2 for t in tot] + [t - t//2 for t in tot]
else:
pads = [t - t//2 for t in tot] + [t//2 for t in tot]
n.attribute.remove(ap[0])
n.attribute.append(helper.make_attribute("pads", pads))
n_fixed += 1
print(" fixed", n.op_type, n.name, "pads", pads)
# 2) optionally fix input dims
if mode.startswith("static"):
hw = [int(x) for x in mode.split(":")[1].split(",")]
x = m.graph.input[0]
d = x.type.tensor_type.shape.dim
d[0].dim_value = 1; d[0].ClearField("dim_param")
if not d[2].dim_value: d[2].dim_value = hw[0]; d[2].ClearField("dim_param")
d[3].dim_value = hw[1]; d[3].ClearField("dim_param")
# clear output shapes so inference recomputes
for o in m.graph.output:
o.type.tensor_type.ClearField("shape")
del m.graph.value_info[:]
m = onnx.shape_inference.infer_shapes(m)
o = m.graph.output[0]
if not o.type.tensor_type.HasField("shape") or len(o.type.tensor_type.shape.dim)==0:
# data-dependent graph (Shape/Reshape): set the known shape explicitly
C = int(sys.argv[4]) if len(sys.argv) > 4 else None
shp = [1, 1, hw[0], hw[1]] if C is None else [1, hw[1]//8, C]
o.type.CopyFrom(helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, shp))
print(" output shape forced to", shp)
print(" input", [dd.dim_value or dd.dim_param for dd in m.graph.input[0].type.tensor_type.shape.dim],
"output", [dd.dim_value or dd.dim_param for dd in m.graph.output[0].type.tensor_type.shape.dim])
onnx.checker.check_model(m)
onnx.save(m, dst)
print(dst, "auto_pad fixed:", n_fixed)
|