如何在ONNX Runtime的C++环境中实现GPU反向传播？

核心前提 #

要实现反向传播，你必须使用ONNX Runtime Training模块（而非普通的推理模块），该模块专门支持训练流程，包括梯度计算、参数更新等操作。

第一步：PyTorch导出支持训练的ONNX模型 #

普通导出的ONNX仅包含前向计算节点，无法直接用于反向传播。需按以下方式导出带训练信息的模型：

1. 确保模型处于训练模式
   导出前调用model.train()，避免PyTorch自动启用推理优化（如Dropout/BatchNorm的推理模式）。

2. 配置ONNX导出参数
   使用torch.onnx.export时需设置关键参数：

   • export_params=True：导出模型参数（必须，因为训练需要更新参数）
   • do_constant_folding=False：禁用常量折叠，保留可训练参数的计算路径
   • training=torch.onnx.TrainingMode.TRAINING：明确导出训练模式的模型
   • input_names/output_names：指定输入输出节点名称，方便C++端定位
   • dynamic_axes：如果需要动态输入尺寸，需配置该参数

3. 示例导出代码

import torch
import torch.nn as nn

# 定义简单的可训练模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
    
    def forward(self, x):
        return self.linear(x)

model = SimpleModel()
model.train()

# 准备示例输入
dummy_input = torch.randn(1, 10)

# 导出ONNX模型
torch.onnx.export(
    model,
    dummy_input,
    "trainable_model.onnx",
    export_params=True,
    do_constant_folding=False,
    training=torch.onnx.TrainingMode.TRAINING,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

第二步：C++端基于ONNX Runtime Training实现GPU反向传播 #

环境准备 #

• 编译/安装带Training模块和CUDA支持的ONNX Runtime：
  官方预编译包可能不包含Training模块，需自行编译，编译时开启ORT_ENABLE_TRAINING和ORT_ENABLE_CUDA选项；或使用支持Training的预编译包（需确认版本说明）。
• 确保系统CUDA、cuDNN版本与ONNX Runtime的CUDA版本匹配。

C++核心实现步骤 #

1. 初始化训练会话与CUDA设备

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

int main() {
    // 配置CUDA执行提供者
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "TrainingSession");
    Ort::SessionOptions session_options;
    
    // 启用CUDA，设置设备ID为0
    OrtCUDAProviderOptions cuda_options;
    cuda_options.device_id = 0;
    session_options.AppendExecutionProvider_CUDA(cuda_options);
    
    // 初始化训练会话，加载ONNX模型
    Ort::TrainingSession training_session(env, "trainable_model.onnx", session_options);

2. 准备GPU输入数据
   需将数据从CPU拷贝到GPU内存，使用Ort::MemoryInfo指定CUDA设备：

// 定义输入尺寸（示例：batch_size=2, feature_dim=10）
    const int64_t input_shape[] = {2, 10};
    const size_t input_size = 2 * 10;
    
    // 生成CPU输入数据
    float input_data[20];
    for (int i = 0; i < 20; ++i) input_data[i] = static_cast<float>(rand()) / RAND_MAX;
    
    // 创建CUDA内存信息
    Ort::MemoryInfo memory_info("Cuda", OrtDeviceAllocator, 0, OrtMemTypeDefault);
    
    // 将数据拷贝到GPU张量
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_data, input_size, input_shape, 2
    );

3. 前向传播、损失计算与反向传播

// 前向传播
    const char* input_names[] = {"input"};
    const char* output_names[] = {"output"};
    Ort::Value output_tensor = training_session.Run(
        Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, 1
    );
    
    // 模拟损失计算（此处需根据你的任务自定义损失，示例用L2损失）
    float* output_data = output_tensor.GetMutableTensorData<float>();
    float target_data[2] = {0.0f, 1.0f}; // 示例目标值
    float loss = 0.0f;
    for (int i = 0; i < 2; ++i) {
        float diff = output_data[i] - target_data[i];
        loss += diff * diff;
    }
    loss /= 2;
    
    // 反向传播：需将损失相关的梯度传入，此处直接基于输出计算梯度并反向传递
    // 注意：实际场景中需根据损失函数构造输出梯度张量
    float output_grad_data[2];
    for (int i = 0; i < 2; ++i) {
        output_grad_data[i] = output_data[i] - target_data[i];
    }
    Ort::Value output_grad_tensor = Ort::Value::CreateTensor<float>(
        memory_info, output_grad_data, 2, input_shape, 1 // 输出尺寸为[2]
    );
    
    const char* output_grad_names[] = {"output"};
    training_session.RunBackward(
        Ort::RunOptions{nullptr}, output_grad_names, &output_grad_tensor, 1
    );

4. 参数更新（使用内置优化器）
   ONNX Runtime Training支持内置优化器（如SGD、Adam），可通过配置启用：

// 初始化SGD优化器
    Ort::TrainingSession::OptimizerOptions optimizer_options;
    optimizer_options.learning_rate = 0.01f;
    training_session.InitOptimizer(optimizer_options);
    
    // 更新参数
    training_session.Optimize();
    
    // 重置梯度，准备下一轮迭代
    training_session.LazyResetGrad();
    
    return 0;
}

关键注意事项 #

• 并非所有PyTorch操作都支持导出为带反向节点的ONNX：如果模型使用自定义操作，需确保该操作的ONNX定义包含反向传播逻辑，或替换为PyTorch内置的、ONNX支持训练的操作。
• 模型参数的维度与类型必须前后一致：PyTorch导出时的参数类型需与C端使用的类型匹配（如float32对应C的float）。
• GPU内存管理：确保输入、输出、梯度张量都分配在CUDA设备上，避免跨设备数据拷贝导致的性能损耗。

内容的提问来源于stack exchange，提问作者Haris Naveed