ilgeco
207 lines
7.1 KiB
Python
207 lines
7.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate output-only ONNX variants for the real yolo11n depth_35 model."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import onnx
|
|
from onnx import TensorProto, helper, shape_inference
|
|
|
|
|
|
ORIGINAL_DEBUG_OUTPUTS = [
|
|
"/model.20/act/Mul_output_0",
|
|
"/model.9/cv1/act/Mul_output_0",
|
|
"/model.8/m.0/m/m.0/cv2/conv/Conv_output_0",
|
|
"/model.23/cv3.1/cv3.1.0/cv3.1.0.0/act/Mul_output_0",
|
|
"/model.8/m.0/m/m.1/cv1/act/Mul_output_0",
|
|
"/model.23/dfl/Transpose_output_0",
|
|
"output0",
|
|
"/model.23/cv2.1/cv2.1.0/act/Mul_output_0",
|
|
]
|
|
|
|
LOCALIZATION_NODE_NAMES = [
|
|
"/model.23/dfl/Transpose",
|
|
"/model.23/dfl/Softmax",
|
|
"/model.23/dfl/conv/Conv",
|
|
"/model.23/dfl/Reshape_1",
|
|
"/model.23/Slice",
|
|
"/model.23/Slice_1",
|
|
"/model.23/Sub",
|
|
"/model.23/Add_1",
|
|
"/model.23/Sub_1",
|
|
"/model.23/Concat_4",
|
|
"/model.23/Mul_2",
|
|
"/model.23/Sigmoid",
|
|
"/model.23/Concat_5",
|
|
]
|
|
|
|
|
|
def collect_value_infos(model: onnx.ModelProto) -> dict[str, onnx.ValueInfoProto]:
|
|
infos: dict[str, onnx.ValueInfoProto] = {}
|
|
for value in list(model.graph.input) + list(model.graph.output) + list(model.graph.value_info):
|
|
infos[value.name] = value
|
|
return infos
|
|
|
|
|
|
def clone_value_info(value: onnx.ValueInfoProto, name: str | None = None) -> onnx.ValueInfoProto:
|
|
cloned = onnx.ValueInfoProto()
|
|
cloned.CopyFrom(value)
|
|
if name is not None:
|
|
cloned.name = name
|
|
return cloned
|
|
|
|
|
|
def make_tensor_value_info_from_type(name: str, tensor_type: onnx.TypeProto.Tensor) -> onnx.ValueInfoProto:
|
|
shape = []
|
|
for dim in tensor_type.shape.dim:
|
|
if dim.HasField("dim_value"):
|
|
shape.append(dim.dim_value)
|
|
elif dim.HasField("dim_param"):
|
|
shape.append(dim.dim_param)
|
|
else:
|
|
shape.append(None)
|
|
return helper.make_tensor_value_info(name, tensor_type.elem_type, shape)
|
|
|
|
|
|
def lookup_output_value_info(model: onnx.ModelProto, value_infos: dict[str, onnx.ValueInfoProto], output_name: str) -> onnx.ValueInfoProto:
|
|
value = value_infos.get(output_name)
|
|
if value is not None:
|
|
return clone_value_info(value)
|
|
|
|
for initializer in model.graph.initializer:
|
|
if initializer.name != output_name:
|
|
continue
|
|
dims = list(initializer.dims)
|
|
return helper.make_tensor_value_info(output_name, initializer.data_type, dims)
|
|
|
|
raise KeyError(f"missing value info for output {output_name}")
|
|
|
|
|
|
def build_model_with_outputs(
|
|
base_model: onnx.ModelProto,
|
|
inferred_model: onnx.ModelProto,
|
|
output_names: list[str],
|
|
extra_nodes: list[onnx.NodeProto] | None = None,
|
|
extra_value_infos: list[onnx.ValueInfoProto] | None = None,
|
|
) -> onnx.ModelProto:
|
|
value_infos = collect_value_infos(inferred_model)
|
|
if extra_value_infos:
|
|
for value in extra_value_infos:
|
|
value_infos[value.name] = value
|
|
model = onnx.ModelProto()
|
|
model.CopyFrom(base_model)
|
|
|
|
del model.graph.output[:]
|
|
for output_name in output_names:
|
|
model.graph.output.append(lookup_output_value_info(inferred_model, value_infos, output_name))
|
|
|
|
if extra_nodes:
|
|
model.graph.node.extend(extra_nodes)
|
|
if extra_value_infos:
|
|
model.graph.value_info.extend(extra_value_infos)
|
|
|
|
onnx.checker.check_model(model)
|
|
return model
|
|
|
|
|
|
def find_node_output(model: onnx.ModelProto, node_name: str) -> str:
|
|
for node in model.graph.node:
|
|
if node.name == node_name:
|
|
return node.output[0]
|
|
|
|
matching_names = sorted(node.name for node in model.graph.node if "model.23" in node.name and "dfl" in node.name)
|
|
suffix = ""
|
|
if matching_names:
|
|
suffix = "\nmatching /model.23 dfl nodes:\n " + "\n ".join(matching_names)
|
|
raise KeyError(f"could not find node named {node_name}{suffix}")
|
|
|
|
|
|
def save_variant(model: onnx.ModelProto, out_dir: Path, variant_name: str) -> None:
|
|
variant_dir = out_dir / variant_name
|
|
variant_dir.mkdir(parents=True, exist_ok=True)
|
|
onnx.save(model, variant_dir / f"{variant_name}.onnx")
|
|
|
|
|
|
def unique_preserving_order(names: list[str]) -> list[str]:
|
|
seen: set[str] = set()
|
|
unique_names: list[str] = []
|
|
for name in names:
|
|
if name in seen:
|
|
continue
|
|
seen.add(name)
|
|
unique_names.append(name)
|
|
return unique_names
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--input", required=True, help="Path to yolo11n depth_35 ONNX model.")
|
|
parser.add_argument("--out-dir", required=True, help="Directory where variants will be generated.")
|
|
args = parser.parse_args()
|
|
|
|
input_path = Path(args.input).resolve()
|
|
out_dir = Path(args.out_dir).resolve()
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
base_model = onnx.load(input_path)
|
|
inferred_model = shape_inference.infer_shapes(base_model)
|
|
|
|
output0_only = build_model_with_outputs(base_model, inferred_model, ["output0"])
|
|
save_variant(output0_only, out_dir, "output0_only")
|
|
|
|
output0_first = build_model_with_outputs(
|
|
base_model,
|
|
inferred_model,
|
|
["output0"] + [name for name in ORIGINAL_DEBUG_OUTPUTS if name != "output0"],
|
|
)
|
|
save_variant(output0_first, out_dir, "output0_first_with_original_debug_outputs")
|
|
|
|
output0_last = build_model_with_outputs(
|
|
base_model,
|
|
inferred_model,
|
|
[name for name in ORIGINAL_DEBUG_OUTPUTS if name != "output0"] + ["output0"],
|
|
)
|
|
save_variant(output0_last, out_dir, "output0_last_with_original_debug_outputs")
|
|
|
|
identity_name = "output0_identity"
|
|
identity_node = helper.make_node("Identity", ["output0"], [identity_name], name="output0_identity_node")
|
|
output0_value = lookup_output_value_info(inferred_model, collect_value_infos(inferred_model), "output0")
|
|
duplicated = build_model_with_outputs(
|
|
base_model,
|
|
inferred_model,
|
|
["output0", identity_name],
|
|
extra_nodes=[identity_node],
|
|
extra_value_infos=[make_tensor_value_info_from_type(identity_name, output0_value.type.tensor_type)],
|
|
)
|
|
save_variant(duplicated, out_dir, "output0_duplicated")
|
|
|
|
localization_outputs = [
|
|
"/model.23/dfl/Transpose_output_0",
|
|
find_node_output(base_model, "/model.23/dfl/Softmax"),
|
|
find_node_output(base_model, "/model.23/dfl/conv/Conv"),
|
|
find_node_output(base_model, "/model.23/dfl/Reshape_1"),
|
|
find_node_output(base_model, "/model.23/Slice"),
|
|
find_node_output(base_model, "/model.23/Slice_1"),
|
|
find_node_output(base_model, "/model.23/Sub"),
|
|
find_node_output(base_model, "/model.23/Add_1"),
|
|
find_node_output(base_model, "/model.23/Sub_1"),
|
|
find_node_output(base_model, "/model.23/Concat_4"),
|
|
find_node_output(base_model, "/model.23/Mul_2"),
|
|
find_node_output(base_model, "/model.23/Sigmoid"),
|
|
find_node_output(base_model, "/model.23/Concat_5"),
|
|
"output0",
|
|
]
|
|
localization = build_model_with_outputs(
|
|
base_model,
|
|
inferred_model,
|
|
unique_preserving_order(localization_outputs),
|
|
)
|
|
save_variant(localization, out_dir, "yolo_tail_localization_outputs")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|