ONNXモデルに出力ノードを追加する方法

ONNXモデルの途中層から出力を得たかったので調べた結果をメモります。

import onnx

# モデルの出力ファイル名
model_path = "path to model"

model = onnx.load(model_path)

input_path = 'path to load model'
output_path = 'path to save model'
input_names = ['input_1']
output_names = ['Logits_softmax', 'Logits_softmax_Flatten'] # 左が元々あったノード　、　右が追加するノード

onnx.utils.extract_model(input_path, output_path, input_names, output_names)

ピンクの丸が追加した出力ノードになります。
これでSoftmaxの手前の出力が取得できるようになりました。