Export a PyTorch Model to ONNX and Extract Intermediate Layer Features

Background : In production projects, extracting intermediate layer features from a model is often required for debugging mismatched inference results, modular deployment, feature transfer, clustering analysis, and model interpretability.

ONNX Overview : ONNX (Open Neural Network Exchange) is an open model format that enables model transfer across frameworks such as PyTorch, TensorFlow, and MXNet, and supports deployment on back‑ends like TensorRT, OpenVINO, ONNX Runtime, and browsers. Its main advantages are cross‑platform deployment, debuggability (graph inspection), and acceleration support.

Example Model : The tutorial uses the image‑classification model mambaout_kobe from the HuggingFace timm repository ( https://huggingface.co/timm/mambaout_kobe.in1k). The model is loaded, set to evaluation mode, and a dummy input tensor is created for tracing.

import timm
import torch
model = timm.create_model('mambaout_kobe', pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "mambaout_kobe.onnx",
    opset_version=11
)

Inspecting ONNX Nodes : Using the onnx Python API, the script iterates over model.graph.node and prints each node’s output names. The target intermediate node is identified as /head/pre_logits/norm/Add_1_output_0.

import onnx
model_path = 'mambaout_kobe.onnx'
model = onnx.load(model_path)
for node in model.graph.node:
    print(node.output)

Adding the Intermediate Output : A ValueInfoProto object describes the desired hidden‑layer output. Its name, data type, and shape (1152) are set, then the object is appended to model.graph.output and the modified model is saved.

from onnx import helper, TensorProto

intermediate_output = helper.ValueInfoProto()
intermediate_output.name = '/head/pre_logits/norm/Add_1_output_0'
intermediate_output.type.tensor_type.elem_type = TensorProto.FLOAT
intermediate_output.type.tensor_type.shape.dim.add().dim_value = 1152

model.graph.output.append(intermediate_output)
onnx.save(model, 'mambaout_kobe_hidden.onnx')

Running with ONNX Runtime : The updated ONNX model is loaded with ort.InferenceSession. Inference returns two outputs: the classification result and the extracted intermediate feature.

sess = ort.InferenceSession("mambaout_kobe_hidden.onnx")
outputs = sess.run(None, {"input": img})
print("Classification shape:", outputs[0].shape)
print("Intermediate feature shape:", outputs[1].shape)

Conclusion : The complete workflow—from PyTorch model loading, ONNX export, node inspection, output modification, to runtime inference—demonstrates how to obtain hidden‑layer features for error localization, modular deployment, and feature analysis in real projects.