Skip to content
编程语言书库

第27章 案例4：机器学习模型推理

机器学习模型部署是从研究到生产的关键环节。本章通过一个完整的 ONNX Runtime 推理封装案例，展示如何高效地将 C++ 推理引擎暴露给 Python，特别关注 GPU 加速和批量处理。

27.1 ONNX Runtime 集成

Section titled “27.1 ONNX Runtime 集成”

ONNX Runtime 简介

Section titled “ONNX Runtime 简介”
┌─────────────────────────────────────────────────────────────────────────────┐
│                         ONNX Runtime 架构                                   │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│   Python API                    C++ API                                     │
│        │                            │                                       │
│        ▼                            ▼                                       │
│   ┌─────────┐                 ┌─────────────┐                              │
│   │ Session │                 │  Inference  │                              │
│   │  Run    │                 │    Engine   │                              │
│   └─────────┘                 └─────────────┘                              │
│        │                            │                                       │
│        ▼                            ▼                                       │
│   ┌─────────────────────────────────────────┐                              │
│   │           Execution Providers           │                              │
│   │  ┌────────┐ ┌────────┐ ┌────────┐       │                              │
│   │  │  CUDA  │ │  ROCm  │ │  CPU   │ ...    │                              │
│   │  └────────┘ └────────┘ └────────┘       │                              │
│   └─────────────────────────────────────────┘                              │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘


关键优势：
- 统一的推理接口
- 多后端支持 (CUDA, ROCm, TensorRT, CoreML)
-  自动优化 (图优化, 算子融合)
-  跨平台部署

环境配置

Section titled “环境配置”
onnx_module.cpp
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <onnxruntime_cxx_api.h>


namespace py = pybind11;


class ONNXModel {
public:
    ONNXModel() : session_(nullptr), env_(nullptr) {}


    ~ONNXModel() {
        if (session_) delete session_;
        // env_ 是单例，通常不删除
    }


    // 加载模型
    void load(const std::string& model_path,
              const std::string& provider = "CPUExecutionProvider") {
        // 创建环境 (单例)
        if (!env_) {
            std::lock_guard<std::mutex> lock(mutex_);
            if (!env_) {
                env_ = new Ort::Env(ORT_LOGGING_LEVEL_WARNING, "ONNXModel");
            }
        }


        // 会话选项
        Ort::SessionOptions session_options;
        session_options.SetIntraOpNumThreads(4);
        session_options.SetGraphOptimizationLevel(
            GraphOptimizationLevel::ORT_ENABLE_ALL
        );


        // 注册执行Provider
        if (provider == "CUDAExecutionProvider") {
            OrtCUDAProviderOptions cuda_options;
            cuda_options.device_id = 0;
            cuda_options.cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::DEFAULT;
            cuda_options.gpu_mem_limit = 0;  // 无限制
            session_options.AppendExecutionProvider_CUDA(cuda_options);
        }


        // 创建会话
        session_ = new Ort::Session(*env_, model_path.c_str(), session_options);


        // 获取输入/输出信息
        size_t num_input_nodes = session_->GetInputCount();
        input_names_.resize(num_input_nodes);
        inputShapes_.resize(num_input_nodes);


        for (size_t i = 0; i < num_input_nodes; ++i) {
            auto input_name = session_->GetInputNameAllocated(i, allocator_);
            input_names_[i] = input_name.get();


            auto type_info = session_->GetInputTypeInfo(i);
            auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
            inputShapes_[i] = tensor_info.GetShape();
        }


        size_t num_output_nodes = session_->GetOutputCount();
        output_names_.resize(num_output_nodes);


        for (size_t i = 0; i < num_output_nodes; ++i) {
            auto output_name = session_->GetOutputNameAllocated(i, allocator_);
            output_names_[i] = output_name.get();
        }
    }


private:
    Ort::Env* env_;
    Ort::Session* session_;
    std::vector<const char*> input_names_;
    std::vector<const char*> output_names_;
    std::vector<std::vector<int64_t>> inputShapes_;
    Ort::AllocatorWithDefaultOptions allocator_;
    static std::mutex mutex_;
};


std::mutex ONNXModel::mutex_;

模型加载接口

Section titled “模型加载接口”
model_loader.cpp
#include <pybind11/pybind11.h>
#include <onnxruntime_cxx_api.h>


namespace py = pybind11;


// 工厂函数
std::unique_ptr<ONNXModel> create_model() {
    return std::make_unique<ONNXModel>();
}


void load_model(ONNXModel& model,
                const std::string& model_path,
                const std::string& provider = "CPUExecutionProvider") {
    model.load(model_path, provider);
}


// 获取模型信息
py::dict get_model_info(ONNXModel& model) {
    py::dict info;


    info["input_names"] = model.get_input_names();
    info["input_shapes"] = model.get_input_shapes();
    info["output_names"] = model.get_output_names();


    return info;
}


py::list get_output_shapes(ONNXModel& model) {
    py::list shapes;
    for (const auto& shape : model.get_output_shapes()) {
        py::list dims;
        for (int64_t d : shape) {
            dims.append(d);
        }
        shapes.append(dims);
    }
    return shapes;
}

27.2 模型推理

Section titled “27.2 模型推理”

核心推理实现

Section titled “核心推理实现”
inference.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <onnxruntime_cxx_api.h>


namespace py = pybind11;


// 推理核心函数
py::array_t<float> infer(
    ONNXModel& model,
    const py::array_t<float>& input_data,
    const std::string& input_name = "") {


    auto input_buf = input_data.request();


    // 创建输入 tensor
    std::vector<int64_t> input_shape;
    for (size_t i = 0; i < input_buf.ndim; ++i) {
        input_shape.push_back(input_buf.shape[i]);
    }


    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        *model.get_allocator(),
        input_shape.data(), input_shape.size(),
        static_cast<float*>(input_buf.ptr),
        input_buf.size * sizeof(float)
    );


    // 运行推理
    const char* input_name_ptr = input_name.empty() ?
        model.get_input_names()[0] : input_name.c_str();


    auto output_tensors = model.run(
        &input_name_ptr, &input_tensor, 1,
        model.get_output_names().data(), model.get_output_names().size()
    );


    // 获取输出
    auto& output_tensor = output_tensors[0];
    auto output_info = output_tensor.GetTensorTypeAndShapeInfo();
    auto output_shape = output_info.GetShape();
    auto output_data = output_tensor.GetTensorMutableData<float>();


    // 转换为 numpy 数组
    std::vector<py::ssize_t> shape;
    for (int64_t d : output_shape) {
        shape.push_back(d);
    }


    // 计算总元素数
    size_t total = 1;
    for (int64_t d : output_shape) {
        total *= d;
    }


    return py::array_t<float>(
        shape, {sizeof(float)}, output_data
    );
}

多输入输出处理

Section titled “多输入输出处理”
multi_io.cpp
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <onnxruntime_cxx_api.h>


namespace py = pybind11;


// 多输入推理
py::list infer_multi_inputs(
    ONNXModel& model,
    const std::vector<py::array_t<float>>& inputs,
    const std::vector<std::string>& input_names = {}) {


    // 创建输入 tensors
    std::vector<Ort::Value> input_tensors;
    std::vector<const char*> input_names_ptr;


    for (size_t i = 0; i < inputs.size(); ++i) {
        const auto& input = inputs[i];
        auto buf = input.request();


        std::vector<int64_t> shape;
        for (size_t j = 0; j < buf.ndim; ++j) {
            shape.push_back(buf.shape[j]);
        }


        const char* name = input_names.empty() ?
            model.get_input_names()[i] : input_names[i].c_str();
        input_names_ptr.push_back(name);


        input_tensors.push_back(Ort::Value::CreateTensor<float>(
            *model.get_allocator(),
            shape.data(), shape.size(),
            static_cast<float*>(buf.ptr),
            buf.size * sizeof(float)
        ));
    }


    // 执行推理
    std::vector<const char*> output_names_ptr;
    for (const auto& name : model.get_output_names()) {
        output_names_ptr.push_back(name);
    }


    auto output_tensors = model.run(
        input_names_ptr.data(), input_tensors.data(), input_tensors.size(),
        output_names_ptr.data(), output_names_ptr.size()
    );


    // 转换为 Python 列表
    py::list results;
    for (auto& tensor : output_tensors) {
        auto tensor_info = tensor.GetTensorTypeAndShapeInfo();
        auto shape = tensor_info.GetShape();
        auto data = tensor.GetTensorMutableData<float>();


        std::vector<py::ssize_t> py_shape;
        for (int64_t d : shape) py_shape.push_back(d);


        size_t total = 1;
        for (int64_t d : shape) total *= d;


        results.append(py::array_t<float>(py_shape, {sizeof(float)}, data));
    }


    return results;
}


// 批量推理
py::list infer_batch(
    ONNXModel& model,
    const py::array_t<float>& batch_input,
    const std::string& input_name = "") {


    auto buf = batch_input.request();


    // 假设第一维是 batch
    std::vector<int64_t> shape;
    for (size_t i = 0; i < buf.ndim; ++i) {
        shape.push_back(buf.shape[i]);
    }


    const char* name = input_name.empty() ?
        model.get_input_names()[0] : input_name.c_str();


    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        *model.get_allocator(),
        shape.data(), shape.size(),
        static_cast<float*>(buf.ptr),
        buf.size * sizeof(float)
    );


    auto output_tensors = model.run(
        &name, &input_tensor, 1,
        model.get_output_names().data(), model.get_output_names().size()
    );


    py::list results;
    for (auto& tensor : output_tensors) {
        auto tensor_info = tensor.GetTensorTypeAndShapeInfo();
        auto out_shape = tensor_info.GetShape();
        auto data = tensor.GetTensorMutableData<float>();


        std::vector<py::ssize_t> py_shape;
        for (int64_t d : out_shape) py_shape.push_back(d);


        results.append(py::array_t<float>(py_shape, {sizeof(float)}, data));
    }


    return results;
}

27.3 输入输出处理

Section titled “27.3 输入输出处理”

Tensor 格式转换

Section titled “Tensor 格式转换”
tensor_utils.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <algorithm>


namespace py = pybind11;


// 类型映射
enum class DataType {
    FLOAT32,
    FLOAT64,
    INT32,
    INT64,
    UINT8
};


py::dtype get_numpy_dtype(DataType dt) {
    switch (dt) {
        case DataType::FLOAT32: return py::dtype::of<float>();
        case DataType::FLOAT64: return py::dtype::of<double>();
        case DataType::INT32: return py::dtype::of<int32_t>();
        case DataType::INT64: return py::dtype::of<int64_t>();
        case DataType::UINT8: return py::dtype::of<uint8_t>();
        default: return py::dtype::of<float>();
    }
}


// NumPy -> ONNX Tensor
template <typename T>
Ort::Value numpy_to_tensor(
    const py::array_t<T>& arr,
    Ort::Allocator& allocator) {


    auto buf = arr.request();


    std::vector<int64_t> shape;
    for (size_t i = 0; i < buf.ndim; ++i) {
        shape.push_back(buf.shape[i]);
    }


    return Ort::Value::CreateTensor<T>(
        allocator, shape.data(), shape.size(),
        static_cast<T*>(buf.ptr), buf.size
    );
}


// ONNX Tensor -> NumPy
template <typename T>
py::array_t<T> tensor_to_numpy(
    const Ort::Value& tensor,
    const std::vector<py::ssize_t>& shape,
    T* data) {


    std::vector<py::ssize_t> strides(shape.size());
    strides.back() = sizeof(T);
    for (int i = static_cast<int>(shape.size()) - 2; i >= 0; --i) {
        strides[i] = strides[i + 1] * shape[i + 1];
    }


    return py::array_t<T>(shape, strides.data(), data);
}

形状验证与调整

Section titled “形状验证与调整”
shape_utils.cpp
#include <pybind11/pybind11.h>
#include <vector>
#include <string>


namespace py = pybind11;


// 形状兼容性检查
bool check_shape_compatible(
    const std::vector<int64_t>& actual,
    const std::vector<int64_t>& expected) {


    if (actual.size() != expected.size()) return false;


    for (size_t i = 0; i < actual.size(); ++i) {
        // -1 表示动态维度
        if (expected[i] != -1 && actual[i] != expected[i]) {
            return false;
        }
    }
    return true;
}


// 自动调整形状
std::vector<int64_t> reshape_for_batch(
    const std::vector<int64_t>& shape,
    int batch_size) {


    std::vector<int64_t> new_shape = shape;


    if (new_shape.empty()) {
        new_shape.push_back(batch_size);
    } else {
        new_shape[0] = batch_size;
    }


    return new_shape;
}


// 广播处理
py::array_t<float> broadcast_arrays(
    const py::array_t<float>& a,
    const py::array_t<float>& b) {


    auto a_buf = a.request();
    auto b_buf = b.request();


    // 计算广播后的形状
    size_t max_ndim = std::max(a_buf.ndim, b_buf.ndim);
    std::vector<py::ssize_t> out_shape(max_ndim);


    for (size_t i = 0; i < max_ndim; ++i) {
        size_t a_idx = i < max_ndim - a_buf.ndim ? 0 : i - (max_ndim - a_buf.ndim);
        size_t b_idx = i < max_ndim - b_buf.ndim ? 0 : i - (max_ndim - b_buf.ndim);


        py::ssize_t a_dim = a_buf.shape[a_buf.ndim - 1 - (a_buf.ndim - 1 - a_idx)];
        py::ssize_t b_dim = b_buf.shape[b_buf.ndim - 1 - (b_buf.ndim - 1 - b_idx)];


        out_shape[max_ndim - 1 - i] = std::max(a_dim, b_dim);
    }


    // 计算输出大小
    size_t total = 1;
    for (auto s : out_shape) total *= s;


    py::array_t<float> result(out_shape);
    auto res_buf = result.request();
    float* res_ptr = static_cast<float*>(res_buf.ptr);


    // 填充 (简化版本)
    const float* a_ptr = static_cast<const float*>(a_buf.ptr);
    const float* b_ptr = static_cast<const float*>(b_buf.ptr);


    for (size_t i = 0; i < total; ++i) {
        res_ptr[i] = a_ptr[i % a_buf.size] + b_ptr[i % b_buf.size];
    }


    return result;
}

类型转换

Section titled “类型转换”
type_conversion.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>


namespace py = pybind11;


// 数据类型转换
py::array_t<float> uint8_to_float32(const py::array_t<uint8_t>& input) {
    auto buf = input.request();


    py::array_t<float> output(buf.shape);
    auto out_buf = output.request();


    uint8_t* in_ptr = static_cast<uint8_t*>(buf.ptr);
    float* out_ptr = static_cast<float*>(out_buf.ptr);


    size_t size = buf.size;
    for (size_t i = 0; i < size; ++i) {
        out_ptr[i] = static_cast<float>(in_ptr[i]) / 255.0f;  // 归一化
    }


    return output;
}


py::array_t<uint8_t> float32_to_uint8(const py::array_t<float>& input) {
    auto buf = input.request();


    py::array_t<uint8_t> output(buf.shape);
    auto out_buf = output.request();


    float* in_ptr = static_cast<float*>(buf.ptr);
    uint8_t* out_ptr = static_cast<uint8_t*>(out_buf.ptr);


    size_t size = buf.size;
    for (size_t i = 0; i < size; ++i) {
        out_ptr[i] = static_cast<uint8_t>(std::clamp(in_ptr[i] * 255.0f, 0.0f, 255.0f));
    }


    return output;
}


// 动态类型转换
py::object convert_tensor(
    const py::array_t<float>& input,
    const std::string& target_dtype) {


    if (target_dtype == "float32") {
        return py::cast(input);
    } else if (target_dtype == "float64") {
        py::array_t<double> output(input.request().shape);
        auto in_buf = input.request();
        auto out_buf = output.request();


        float* in_ptr = static_cast<float*>(in_buf.ptr);
        double* out_ptr = static_cast<double*>(out_buf.ptr);


        for (size_t i = 0; i < in_buf.size; ++i) {
            out_ptr[i] = static_cast<double>(in_ptr[i]);
        }


        return py::cast(output);
    } else if (target_dtype == "int32") {
        py::array_t<int32_t> output(input.request().shape);
        auto in_buf = input.request();
        auto out_buf = output.request();


        float* in_ptr = static_cast<float*>(in_buf.ptr);
        int32_t* out_ptr = static_cast<int32_t*>(out_buf.ptr);


        for (size_t i = 0; i < in_buf.size; ++i) {
            out_ptr[i] = static_cast<int32_t>(in_ptr[i]);
        }


        return py::cast(output);
    }


    throw py::value_error("Unsupported target dtype: " + target_dtype);
}

27.4 批量推理优化

Section titled “27.4 批量推理优化”

批量处理策略

Section titled “批量处理策略”
batch_inference.cpp
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <vector>
#include <future>


namespace py = pybind11;


// 动态批量
py::list batch_inference_dynamic(
    ONNXModel& model,
    const std::vector<py::array_t<float>>& inputs,
    size_t max_batch_size = 32) {


    py::list results;


    // 分批处理
    for (size_t i = 0; i < inputs.size(); i += max_batch_size) {
        size_t batch_end = std::min(i + max_batch_size, inputs.size());
        size_t batch_size = batch_end - i;


        // 找到最大形状
        std::vector<size_t> max_dims(inputs[0].request().ndim, 0);
        for (size_t j = i; j < batch_end; ++j) {
            auto buf = inputs[j].request();
            for (size_t d = 0; d < buf.ndim; ++d) {
                max_dims[d] = std::max(max_dims[d], static_cast<size_t>(buf.shape[d]));
            }
        }


        // 创建批次张量
        std::vector<float> batch_data;
        for (size_t j = i; j < batch_end; ++j) {
            auto buf = inputs[j].request();
            float* ptr = static_cast<float*>(buf.ptr);
            batch_data.insert(batch_data.end(), ptr, ptr + buf.size);
        }


        // 构建批次数组
        std::vector<py::ssize_t> batch_shape = {static_cast<py::ssize_t>(batch_size)};
        for (size_t d = 0; d < max_dims.size(); ++d) {
            batch_shape.push_back(max_dims[d]);
        }


        py::array_t<float> batch_input(batch_shape);
        auto batch_buf = batch_input.request();
        std::memcpy(batch_buf.ptr, batch_data.data(), batch_data.size() * sizeof(float));


        // 推理
        auto outputs = infer_batch(model, batch_input);


        // 拆分结果
        for (size_t j = 0; j < batch_size; ++j) {
            results.append(outputs[0]);
        }
    }


    return results;
}


// 预分配缓冲区批量推理
class BatchedInference {
public:
    BatchedInference(ONNXModel& model, size_t max_batch, size_t tensor_size)
        : max_batch_(max_batch), tensor_size_(tensor_size) {


        // 预分配输入缓冲区
        input_buffer_.resize(max_batch * tensor_size);
        input_shape_ = {static_cast<int64_t>(max_batch), 1, 224, 224};


        // 创建预分配的张量
        input_tensor_ = Ort::Value::CreateTensor<float>(
            *allocator_,
            input_shape_.data(), input_shape_.size(),
            input_buffer_.data(), input_buffer_.size()
        );
    }


    py::list run(const std::vector<py::array_t<float>>& inputs) {
        size_t batch_size = std::min(inputs.size(), max_batch_);


        // 填充输入缓冲区
        float* ptr = input_buffer_.data();
        for (size_t i = 0; i < batch_size; ++i) {
            auto buf = inputs[i].request();
            std::memcpy(ptr, buf.ptr, tensor_size_ * sizeof(float));
            ptr += tensor_size_;
        }


        // 更新形状
        input_shape_[0] = batch_size;


        // 推理
        auto outputs = model_->run(...);


        // 处理输出
        py::list results;
        for (size_t i = 0; i < batch_size; ++i) {
            // 提取第 i 个结果
            results.append(outputs[0]);
        }


        return results;
    }


private:
    size_t max_batch_;
    size_t tensor_size_;
    std::vector<float> input_buffer_;
    std::vector<int64_t> input_shape_;
    Ort::Value input_tensor_;
};

并行推理

Section titled “并行推理”
parallel_inference.cpp
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <thread>
#include <future>


namespace py = pybind11;


// 多线程批量处理
py::list parallel_batch_inference(
    ONNXModel& model,
    const std::vector<py::array_t<float>>& inputs,
    size_t num_threads = 4) {


    size_t batch_size = inputs.size();
    size_t batch_per_thread = (batch_size + num_threads - 1) / num_threads;


    std::vector<std::future<py::list>> futures;


    for (size_t t = 0; t < num_threads; ++t) {
        size_t start = t * batch_per_thread;
        size_t end = std::min(start + batch_per_thread, batch_size);


        if (start >= batch_size) break;


        auto future = std::async(std::launch::async, [&](size_t s, size_t e) {
            py::list results;
            for (size_t i = s; i < e; ++i) {
                auto result = infer(model, inputs[i]);
                results.append(result);
            }
            return results;
        }, start, end);


        futures.push_back(std::move(future));
    }


    // 收集结果
    py::list all_results;
    for (auto& f : futures) {
        py::list thread_results = f.get();
        for (py::handle r : thread_results) {
            all_results.append(r);
        }
    }


    return all_results;
}


// 模型池 (用于高并发场景)
class ModelPool {
public:
    ModelPool(const std::string& model_path, size_t pool_size)
        : model_path_(model_path) {


        for (size_t i = 0; i < pool_size; ++i) {
            auto model = std::make_unique<ONNXModel>();
            model->load(model_path);
            pool_.push(std::move(model));
        }
    }


    std::shared_ptr<ONNXModel> acquire() {
        std::unique_lock<std::mutex> lock(mutex_);
        cv_.wait(lock, [this] { return !pool_.empty(); });


        auto model = std::shared_ptr<ONNXModel>(pool_.front().release(), [&](ONNXModel* m) {
            std::lock_guard<std::mutex> lock(mutex_);
            pool_.push(std::unique_ptr<ONNXModel>(m));
            cv_.notify_one();
        });


        return model;
    }


private:
    std::string model_path_;
    std::queue<std::unique_ptr<ONNXModel>> pool_;
    std::mutex mutex_;
    std::condition_variable cv_;
};

27.5 GPU 加速

Section titled “27.5 GPU 加速”

CUDA 内存管理

Section titled “CUDA 内存管理”
cuda_utils.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <cuda_runtime.h>


namespace py = pybind11;


// CUDA 内存包装
class CUDABuffer {
public:
    CUDABuffer() : ptr_(nullptr), size_(0) {}


    CUDABuffer(size_t size) : size_(size) {
        cudaMalloc(&ptr_, size);
    }


    ~CUDABuffer() {
        if (ptr_) cudaFree(ptr_);
    }


    void* ptr() { return ptr_; }
    size_t size() const { return size_; }


    // 从 CPU 复制到 GPU
    void upload(const void* cpu_ptr, size_t size) {
        if (size > size_) {
            cudaFree(ptr_);
            cudaMalloc(&ptr_, size);
            size_ = size;
        }
        cudaMemcpy(ptr_, cpu_ptr, size, cudaMemcpyHostToDevice);
    }


    // 从 GPU 复制到 CPU
    void download(void* cpu_ptr, size_t size) const {
        cudaMemcpy(cpu_ptr, ptr_, size, cudaMemcpyDeviceToHost);
    }


    // 移动语义
    CUDABuffer(CUDABuffer&& other) noexcept
        : ptr_(other.ptr_), size_(other.size_) {
        other.ptr_ = nullptr;
        other.size_ = 0;
    }


    CUDABuffer& operator=(CUDABuffer&& other) noexcept {
        if (this != &other) {
            if (ptr_) cudaFree(ptr_);
            ptr_ = other.ptr_;
            size_ = other.size_;
            other.ptr_ = nullptr;
            other.size_ = 0;
        }
        return *this;
    }


private:
    void* ptr_;
    size_t size_;
};


// CUDA 异步操作
class CUDAStream {
public:
    CUDAStream() {
        cudaStreamCreate(&stream_);
    }


    ~CUDAStream() {
        cudaStreamDestroy(stream_);
    }


    void synchronize() {
        cudaStreamSynchronize(stream_);
    }


    cudaStream_t get() { return stream_; }


private:
    cudaStream_t stream_;
};


// GPU 推理
py::array_t<float> infer_gpu(
    ONNXModel& model,
    const py::array_t<float>& input,
    CUDAStream& stream) {


    auto buf = input.request();


    // 创建 GPU 缓冲区
    CUDABuffer gpu_input(buf.size * sizeof(float));
    gpu_input.upload(buf.ptr, buf.size * sizeof(float));


    // 创建 GPU 输出缓冲区
    CUDABuffer gpu_output(buf.size * sizeof(float));


    // 构造输入张量 (使用 GPU 内存)
    std::vector<int64_t> shape;
    for (size_t i = 0; i < buf.ndim; ++i) {
        shape.push_back(buf.shape[i]);
    }


    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        *model.get_allocator(),
        shape.data(), shape.size(),
        gpu_input.ptr(), buf.size * sizeof(float)
    );


    // 运行推理 (GPU)
    auto output_tensors = model.run_gpu(input_tensor, stream.get());


    // 下载结果
    py::array_t<float> result(buf.shape);
    auto res_buf = result.request();
    gpu_output.download(res_buf.ptr, buf.size * sizeof(float));


    stream.synchronize();
    return result;
}

ROCm 支持

Section titled “ROCm 支持”
rocm_utils.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>


namespace py = pybind11;


// ROCm 内存管理 (类似于 CUDA)
class ROCmBuffer {
public:
    ROCmBuffer(size_t size) {
        hipMalloc(&ptr_, size);
        size_ = size;
    }


    ~ROCmBuffer() {
        if (ptr_) hipFree(ptr_);
    }


    void* ptr() { return ptr_; }


    void upload(const void* cpu_ptr, size_t size) {
        hipMemcpy(ptr_, cpu_ptr, size, hipMemcpyHostToDevice);
    }


    void download(void* cpu_ptr, size_t size) const {
        hipMemcpy(cpu_ptr, ptr_, size, hipMemcpyDeviceToHost);
    }


private:
    void* ptr_;
    size_t size_;
};


// ROCm 推理选项
struct ROCmOptions {
    int device_id = 0;
    int gfx_arch = -1;  // 自动检测
    bool use_fp16 = false;
    bool use_fast_math = true;
};


// 配置 ROCm Provider
void configure_rocm_provider(Ort::SessionOptions& options, const ROCmOptions& rocm_opts) {
    OrtROCMProviderOptions rocm_options;
    rocm_options.device_id = rocm_opts.device_id;
    rocm_options.gfx_arch = rocm_opts.gfx_arch;
    rocm_options.fast_math = rocm_opts.use_fast_math;


    options.AppendExecutionProvider_ROCM(rocm_options);
}

内存布局优化

Section titled “内存布局优化”
memory_layout.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <algorithm>


namespace py = pybind11;


// 内存布局转换
// NumPy 默认 row-major (C-order)
// ONNX Runtime 通常期望 batch-first
// GPU 加速可能需要特定布局 (NCHW vs NHWC)


// NCHW -> NHWC 转换 (TensorFlow 风格)
py::array_t<float> nchw_to_nhwc(const py::array_t<float>& input) {
    auto buf = input.request();


    // 假设 [N, C, H, W]
    if (buf.ndim != 4) {
        throw py::value_error("Expected 4D tensor in NCHW format");
    }


    py::ssize_t N = buf.shape[0];
    py::ssize_t C = buf.shape[1];
    py::ssize_t H = buf.shape[2];
    py::ssize_t W = buf.shape[3];


    py::array_t<float> output({N, H, W, C});
    auto out_buf = output.request();


    float* in_ptr = static_cast<float*>(buf.ptr);
    float* out_ptr = static_cast<float*>(out_buf.ptr);


    for (py::ssize_t n = 0; n < N; ++n) {
        for (py::ssize_t h = 0; h < H; ++h) {
            for (py::ssize_t w = 0; w < W; ++w) {
                for (py::ssize_t c = 0; c < C; ++c) {
                    // NCHW: [n, c, h, w] -> NHWC: [n, h, w, c]
                    size_t in_idx = ((n * C + c) * H + h) * W + w;
                    size_t out_idx = ((n * H + h) * W + w) * C + c;
                    out_ptr[out_idx] = in_ptr[in_idx];
                }
            }
        }
    }


    return output;
}


// NHWC -> NCHW 转换
py::array_t<float> nhwc_to_nchw(const py::array_t<float>& input) {
    auto buf = input.request();


    if (buf.ndim != 4) {
        throw py::value_error("Expected 4D tensor in NHWC format");
    }


    py::ssize_t N = buf.shape[0];
    py::ssize_t H = buf.shape[1];
    py::ssize_t W = buf.shape[2];
    py::ssize_t C = buf.shape[3];


    py::array_t<float> output({N, C, H, W});
    auto out_buf = output.request();


    float* in_ptr = static_cast<float*>(buf.ptr);
    float* out_ptr = static_cast<float*>(out_buf.ptr);


    for (py::ssize_t n = 0; n < N; ++n) {
        for (py::ssize_t c = 0; c < C; ++c) {
            for (py::ssize_t h = 0; h < H; ++h) {
                for (py::ssize_t w = 0; w < W; ++w) {
                    size_t in_idx = ((n * H + h) * W + w) * C + c;
                    size_t out_idx = ((n * C + c) * H + h) * W + w;
                    out_ptr[out_idx] = in_ptr[in_idx];
                }
            }
        }
    }


    return output;
}


// 检查内存是否连续
bool is_contiguous(const py::buffer_info& buf) {
    size_t expected_stride = buf.itemsize;


    for (int i = static_cast<int>(buf.ndim) - 1; i >= 0; --i) {
        if (buf.strides[i] != expected_stride) {
            return false;
        }
        expected_stride *= buf.shape[i];
    }


    return true;
}


// 确保内存连续
py::array_t<float> ensure_contiguous(const py::array_t<float>& input) {
    auto buf = input.request();


    if (is_contiguous(buf)) {
        return input;
    }


    py::array_t<float> output(buf.shape);
    auto out_buf = output.request();


    std::memcpy(out_buf.ptr, buf.ptr, buf.size * buf.itemsize);


    return output;
}

Python 接口

Section titled “Python 接口”
import numpy as np
from onnx_wrapper import (
    ONNXModel, infer, infer_batch, infer_multi_inputs,
    ModelPool, CUDABuffer, nchw_to_nhwc
)


model = ONNXModel()
model.load("resnet50.onnx", provider="CUDAExecutionProvider")


input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)


output = infer(model, input_data)
print(f"Output shape: {output.shape}")


batch_input = np.random.randn(8, 3, 224, 224).astype(np.float32)
outputs = infer_batch(model, batch_input)
print(f"Batch output shape: {outputs[0].shape}")


gpu_buffer = CUDABuffer(1024 * 1024 * 1024)  # 1GB
gpu_buffer.upload(input_data.ctypes.data, input_data.nbytes)


pool = ModelPool("resnet50.onnx", pool_size=4)


def parallel_inference(inputs):
    model = pool.acquire()
    try:
        return infer(model, inputs)
    finally:
        pass  # 自动归还


if __name__ == "__main__":
    # 单次推理
    img = np.random.randn(1, 3, 224, 224).astype(np.float32)
    result = infer(model, img)
    print(f"Result: {result.shape}")


    # 批量处理
    images = [np.random.randn(3, 224, 224).astype(np.float32) for _ in range(16)]
    results = infer_batch(model, np.stack(images))
    print(f"Processed {len(results)} images")

关键洞察：ONNX Runtime 是 ML 模型部署的标准选择。通过 pybind11 可以直接将 NumPy 数组传递给推理引擎，无需额外的内存拷贝。GPU 加速的关键在于内存布局匹配：NCHW 用于 CUDA/TensorRT，NHWC 用于 CPU/ROCm。批量推理时使用模型池和预分配缓冲区可以显著提升吞吐量。内存布局转换（行列优先、NCHW/NHWC）往往是性能瓶颈，必须在数据预处理阶段完成。