File size: 3,474 Bytes
6cdca59 | 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 | #!/usr/bin/env python3
"""Remove the unsupported Log branch from this YOLO depth ONNX export.
For the depth head exported by this project:
Log(Clip(Exp(Clip(x, lo, hi)), floor, +inf))
is equivalent to Clip(x, max(lo, log(floor)), hi). The generated AX637
model then contains no AxQuantizedLog CPU fallback.
"""
from __future__ import annotations
import argparse
import math
from pathlib import Path
import numpy as np
import onnx
from onnx import numpy_helper
def scalar_initializer(model: onnx.ModelProto, name: str) -> float:
"""Return a scalar initializer and reject models with a different topology."""
for initializer in model.graph.initializer:
if initializer.name == name:
value = numpy_helper.to_array(initializer)
if value.size != 1:
raise ValueError(f"{name} must be scalar, got shape {value.shape}")
return float(value.reshape(-1)[0])
raise ValueError(f"missing initializer: {name}")
def fold_depth_log(model: onnx.ModelProto) -> None:
"""Fold the Clip -> Exp -> Clip -> Log sequence used by the depth head."""
nodes = list(model.graph.node)
producers = {output: node for node in nodes for output in node.output if output}
log_node = next(
(node for node in nodes if node.name == "/model.23/Log" and node.op_type == "Log"),
None,
)
if log_node is None or len(log_node.input) != 1 or len(log_node.output) != 1:
raise ValueError("expected a single-input /model.23/Log node")
outer_clip = producers.get(log_node.input[0])
exp_node = producers.get(outer_clip.input[0]) if outer_clip is not None and outer_clip.input else None
inner_clip = producers.get(exp_node.input[0]) if exp_node is not None and exp_node.input else None
if (
outer_clip is None
or outer_clip.op_type != "Clip"
or exp_node is None
or exp_node.op_type != "Exp"
or inner_clip is None
or inner_clip.op_type != "Clip"
or len(inner_clip.input) != 3
or len(outer_clip.input) < 2
):
raise ValueError("unexpected depth calibration topology around /model.23/Log")
inner_min = scalar_initializer(model, inner_clip.input[1])
floor = scalar_initializer(model, outer_clip.input[1])
if floor <= 0.0:
raise ValueError(f"Clip floor must be positive, got {floor}")
# Do not mutate a potentially shared initializer from elsewhere in the graph.
folded_min = max(inner_min, math.log(floor))
min_name = "depth_log_folded_min"
model.graph.initializer.append(numpy_helper.from_array(np.array(folded_min, dtype=np.float32), min_name))
inner_clip.input[1] = min_name
inner_clip.name = "/model.23/Clip_log_folded"
inner_clip.output[0] = log_node.output[0]
remove_ids = {id(outer_clip), id(exp_node), id(log_node)}
del model.graph.node[:]
model.graph.node.extend(node for node in nodes if id(node) not in remove_ids)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("input", type=Path, help="source ONNX model")
parser.add_argument("output", type=Path, help="rewritten ONNX model")
args = parser.parse_args()
model = onnx.load(args.input)
fold_depth_log(model)
onnx.checker.check_model(model)
args.output.parent.mkdir(parents=True, exist_ok=True)
onnx.save(model, args.output)
print(f"wrote {args.output}")
if __name__ == "__main__":
main()
|