ONNXRuntime CUDA 版源码编译与 C++ 部署 PyTorch 模型教程

3 Pytorch 模型导出 .onnx

以 SAC 算法的部署为例，核心代码如下所示：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
policy = SACPolicy(
    obs_dim=16,
    action_dim=2,
    device=device,
    **params["policy"],
)
model_path = os.path.abspath(os.path.join(__file__, "../sac_policy/sac_model.pth"))
policy.load_state_dict(torch.load(model_path, map_location=device))

# convert to ONNX model
policy.eval()
dummy_obs = torch.randn(1, policy.obs_dim).to(device)
input_names = ["observation"]
output_names = ["actions"]
output_path = os.path.abspath(os.path.join(__file__, "../sac_model.onnx"))
torch.onnx.export(
    model=policy,
    args=dummy_obs,
    f=output_path,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes={"observation": {0: "batch_size"}, "actions": {0: "batch_size"}},
    verbose=False,
)

可以使用以下脚本验证：

ort_session = ort.InferenceSession(output_path)
dummy_obs_np = dummy_obs.cpu().numpy()
ort_inputs = {ort_session.get_inputs()[0].name: dummy_obs_np}
ort_outputs = ort_session.run(None, ort_inputs)
torch_outputs = policy(dummy_obs)
print(ort_outputs)
for i, (torch_out, ort_out) in enumerate(zip(torch_outputs, ort_outputs)):
    diff = np.abs(torch_out.detach().cpu().numpy() - ort_out)
    print(f"Output {i} 最大差异：{diff.max():.1e}")
print(f"ONNX 模型已成功导出至 {output_path}，并验证通过！")

4 C++ 调用 .onnx

以 SAC 算法的部署为例，核心代码如下所示：

SACPolicy::SACPolicy(const std::string& model_path, bool use_gpu) : env_(ORT_LOGGING_LEVEL_WARNING, "SACPolicyInference") {
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    if (use_gpu) {
        OrtCUDAProviderOptions cuda_options;
        cuda_options.device_id = 0;
        cuda_options.cudnn_conv_algo_search = OrtCudnnConvAlgoSearchExhaustive;
        cuda_options.arena_extend_strategy = 0;
        cuda_options.do_copy_in_default_stream = 0;
        session_options.AppendExecutionProvider_CUDA(cuda_options);
    }
    session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), session_options);
}

SACPredictResult SACPolicy::predict(const std::vector<double>& observation) {
    Ort::AllocatorWithDefaultOptions allocator;
    const auto& input_tensor_info = session_->GetInputTypeInfo(0);
    const auto& input_shape = input_tensor_info.GetTensorTypeAndShapeInfo().GetShape();
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
    
    std::vector<float> input_data(observation.begin(), observation.end());
    std::vector<Ort::Value> input_tensors;
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_data.data(), input_data.size(), input_shape.data(), input_shape.size());
    input_tensors.push_back(std::move(input_tensor));
    
    std::vector<const char*> input_names = {"observation"};
    std::vector<const char*> output_names = {"actions"};
    std::vector<Ort::Value> output_tensors = session_->Run(
        Ort::RunOptions{nullptr},
        input_names.data(),
        input_tensors.data(),
        input_tensors.size(),
        output_names.data(),
        output_names.size()
    );
    
    float* output_ptr = output_tensors[0].GetTensorMutableData<float>();
    SACPredictResult result;
    result.linear_velocity = output_ptr[0];
    result.angular_velocity = output_ptr[1];
    return result;
}