| |
| """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}") |
|
|
| |
| 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() |
|
|
