Get Model Input and Output Shapes using ONNX Runtime and C++

If you are working with ONNX models, it is important to know how to retrieve the input and output shapes of the model. Using this information, you can prepare input data and process your model's outputs. This tutorial shows how to get the model input and output shapes using ONNX Runtime and C++.

In the following code, we load the ONNX model using the Ort::Session class. We retrieve the number of inputs and outputs of the model using the GetInputCount and GetOutputCount functions. Inside the loop, we retrieve the name and shape for each input and output using the GetInputNameAllocated and GetShape functions. Finally, information for each input and output is printed to the console.

#include <iostream>
#include <onnxruntime_cxx_api.h>

std::ostream &operator<<(std::ostream &os, const std::vector<int64_t> &values)
{
    os << "[" << values.front();
    for (int i = 1; i < values.size(); ++i) {
        os << ", " << values[i];
    }
    os << "]";

    return os;
}

int main()
{
    Ort::AllocatorWithDefaultOptions allocator;
    Ort::SessionOptions sessionOptions;
    sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    Ort::Session session(Ort::Env(), "yolov8s.onnx", sessionOptions);

    size_t inputCount = session.GetInputCount();
    for (int i = 0; i < inputCount; ++i) {
        auto name = session.GetInputNameAllocated(i, allocator);
        auto shape = session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();

        std::cout << "Input Number: " << i << std::endl;
        std::cout << " Input Name: " << name.get() << std::endl;
        std::cout << " Input Shape: " << shape << std::endl;
    }

    size_t outputCount = session.GetOutputCount();
    for (int i = 0; i < outputCount; ++i) {
        auto name = session.GetOutputNameAllocated(i, allocator);
        auto shape = session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();

        std::cout << "Output Number: " << i << std::endl;
        std::cout << " Output Name: " << name.get() << std::endl;
        std::cout << " Output Shape: " << shape << std::endl;
    }

    return 0;
}

Here is an example of the console output:

Input Number: 0
 Input Name: images
 Input Shape: [1, 3, 480, 640]
Output Number: 0
 Output Name: output0
 Output Shape: [1, 84, 6300]

In this example, the model has one input which shape is [1, 3, 480, 640] and one output which shape is [1, 84, 6300].