Patching ONNX to expose intermediate output
def patch_onnx(input_filename: str, output_filename: str, inter_layer_names: List[str])->None:
"""
Main idea: append value_info_proto with name included in inter_layer_names to the
graph output list
"""
model = onnx.load(input_filename)
value_info_protos = []
for node in model.graph.value_info:
if node.name in inter_layer_names:
value_info_protos.append(node)
assert len(value_info_protos) == len(inter_layer_names)
model.graph.output.extend(value_info_protos)
onnx.checker.check_model(model)
onnx.save(model, output_filename)
■────────────────────────────────────────────────────────────────────■
│ │
│ Copyright 2014-2024 Aria Ghora Prabono. Any and all opinions │
│ listed here are personal unless stated otherwise. │
│ │
■────────────────────────────────────────────────────────────────────■