3. 创建推理会话并执行单次推理

任务目标

本案例演示如何配置会话选项、创建推理会话,并对随机生成的输入数据执行一次前向推理,最终验证输出张量的形状是否符合预期。案例基于ONNX Runtime C++ API,使用一个公开的ONNX模型文件完成闭环。

环境与依赖

  • 操作系统:Linux / Windows / macOS(支持ONNX Runtime即可)
  • 依赖库:ONNX Runtime(版本≥1.15),CMake(≥3.16),C++编译器(支持C++17)
  • 模型文件:从ONNX Model Zoo下载 mnist-8.onnx,放置于可执行文件同目录

关键实现

以下代码创建会话选项(启用CPU优化),加载模型,构造随机输入(形状1×1×28×28),执行推理并打印输出张量的形状。输出应为1×10,对应10个类别的logits。

#include <onnxruntime_cxx_api.h>
#include <vector>
#include <random>
#include <iostream>
#include <cassert>

int main() {
    // 1. 环境与会话选项
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "session-demo");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    // 2. 创建推理会话(加载模型)
    const char* model_path = "mnist-8.onnx";
    Ort::Session session(env, model_path, session_options);

    // 3. 获取输入输出名称(仅用于打印,非必需)
    Ort::AllocatorWithDefaultOptions allocator;
    auto input_name = session.GetInputNameAllocated(0, allocator);
    auto output_name = session.GetOutputNameAllocated(0, allocator);
    std::cout << "Input: " << input_name.get() << ", Output: " << output_name.get() << std::endl;

    // 4. 构造随机输入数据(MNIST: 1×1×28×28)
    const int64_t input_shape[] = {1, 1, 28, 28};
    const size_t input_size = 1 * 1 * 28 * 28;
    std::vector<float> input_data(input_size);
    std::mt19937 gen(42);  // 固定随机种子,保证可复现
    std::uniform_real_distribution<float> dist(0.0f, 1.0f);
    for (auto& v : input_data) v = dist(gen);

    // 5. 创建输入输出张量
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_data.data(), input_size, input_shape, 4);

    // 6. 执行推理
    const char* input_names[] = {input_name.get()};
    const char* output_names[] = {output_name.get()};
    auto output_tensors = session.Run(Ort::RunOptions{nullptr}, 
                                      input_names, &input_tensor, 1,
                                      output_names, 1);

    // 7. 验证输出形状
    auto& output_tensor = output_tensors.front();
    auto output_shape = output_tensor.GetTensorTypeAndShapeInfo().GetShape();
    std::cout << "Output shape: [";
    for (size_t i = 0; i < output_shape.size(); ++i) {
        std::cout << output_shape[i];
        if (i + 1 < output_shape.size()) std::cout << ", ";
    }
    std::cout << "]" << std::endl;

    // 8. 断言输出形状为 [1, 10]
    assert(output_shape.size() == 2);
    assert(output_shape[0] == 1);
    assert(output_shape[1] == 10);
    std::cout << "Shape verification passed." << std::endl;

    return 0;
}

编译与运行

# CMakeLists.txt(与上述cpp同目录)
cmake_minimum_required(VERSION 3.16)
project(session_demo)
find_package(onnxruntime REQUIRED)
add_executable(session_demo main.cpp)
target_link_libraries(session_demo onnxruntime::onnxruntime)
# 编译与运行
mkdir build && cd build
cmake .. -DONNXRUNTIME_ROOT=/path/to/onnxruntime
make
./session_demo   # 确保 mnist-8.onnx 在当前目录

验证方法

  • 程序正常退出,退出码为0。
  • 控制台输出包含 Output shape: [1, 10] 和 Shape verification passed.。
  • 若断言失败(输出形状非[1,10]),程序将中止并提示断言错误。

预期输出

Input: Input3, Output: Plus214_Output_0
Output shape: [1, 10]
Shape verification passed.

注:输入/输出名称可能因模型版本略有差异,但形状验证结果不变。