Intermediate result inspection in ONNX models

add_debug_output.py

from pathlib import Path
import onnx

model_path = Path("my_model.onnx")
extra_output = "/up1/Concat_3_output_0" # One good way to get this identifier for your node is to use https://netron.app/
extra_output_shape = [1, 1024, 55, 55]
new_model_path = model_path.with_stem(model_path.stem + "-" + extra_output.replace("/", "_").replace(".", "_").strip("_"))

model = onnx.load(model_path)

# Create a new output 
intermediate_layer_value_info = onnx.helper.make_tensor_value_info("debug_out",
                                           onnx.TensorProto.FLOAT,
                                           extra_output_shape)
# Create an identity node to write the intermediate result to the new output
output_identity_node = onnx.helper.make_node("Identity", [extra_output], ["debug_out"])
# Add the output and the new node
model.graph.output.append(intermediate_layer_value_info)
model.graph.node.append(output_identity_node)

# Save the new model
onnx.save(model, new_model_path)