使用TensorRT引擎时如何正确格式化输入与调整输出数据

一、主机到设备数据拷贝（OpenCV → TensorRT） #

1. 输入维度与内存缓冲区的关联 #

TensorRT的内存缓冲区是连续的一维线性存储，维度格式（NCHW/NHWC）决定了多维数据在一维缓冲区中的排列顺序：

• NCHW：按批量数(N) → 通道数(C) → 高度(H) → 宽度(W)的顺序依次存储，每个通道的所有像素连续存放
• NHWC：按批量数(N) → 高度(H) → 宽度(W) → 通道数(C)的顺序依次存储，每个像素的三个通道值连续存放

2. OpenCV NHWC与模型输入NHWC匹配时的处理 #

OpenCV的Mat默认存储格式是NHWC（单张图是HWC，批量则是NHWC），如果TensorRT模型输入也是NHWC格式，且输入尺寸（H、W、C）与OpenCV Mat一致，不需要重新排列缓冲区数据，直接做内存拷贝即可（注意数据类型转换：OpenCV默认是8位无符号char，TensorRT模型输入一般为float）。

3. TensorRT引擎期望的数据格式序列 #

引擎的输入格式由模型导出/转换时的配置决定，可通过代码确认：

// 获取输入绑定的数据格式与张量排列
nvinfer1::DataType input_type = engine->getBindingDataType(input_idx);
nvinfer1::TensorFormat input_format = engine->getBindingTensorFormat(input_idx);

若为NHWC，引擎期望每个像素的通道值连续，按行优先遍历图像；若为NCHW，则每个通道的所有像素连续存放。

4. 适配新输入维度格式的通用方法 #

核心是将OpenCV Mat数据转换成引擎期望的一维缓冲区顺序，步骤如下：

1. 将OpenCV Mat resize到模型要求的H、W尺寸
2. 转换数据类型（如CV_8UC3转CV_32FC3）并按模型要求做归一化
3. 根据引擎的TensorFormat重新排列数据：
   • NHWC格式：直接拷贝Mat的data指针即可
   • NCHW格式：拆分通道后按通道顺序连续存放

针对你给出的NHWC动态输入示例，完整代码如下：

// 假设input_dims是{1, 386, 342, 3}，对应N=1, H=386, W=342, C=3
cv::Mat img = cv::imread("test.jpg");
// 1. resize到模型要求的尺寸（W对应d[2], H对应d[1]）
cv::resize(img, img, cv::Size(input_dims.d[2], input_dims.d[1]));
// 2. 转换数据类型并归一化（根据模型要求调整）
img.convertTo(img, CV_32FC3);
img = img / 255.0f;
// 3. NHWC格式直接拷贝到主机端输入缓冲区
std::memcpy(input_host, img.data, input_size);
// 4. 主机到设备异步拷贝
cudaMemcpyAsync(input_mem, input_host, input_size, cudaMemcpyHostToDevice, stream);

若遇到NCHW格式输入，需调整数据排列：

cv::Mat img_resized;
cv::resize(img, img_resized, cv::Size(input_dims.d[3], input_dims.d[2])); // NCHW的d[2]是H, d[3]是W
img_resized.convertTo(img_resized, CV_32FC3);
img_resized = img_resized / 255.0f;

// 拆分通道并按NCHW顺序排列
std::vector<cv::Mat> channels(3);
cv::split(img_resized, channels);
int channel_size = input_dims.d[2] * input_dims.d[3] * sizeof(float);
std::memcpy(input_host, channels[0].data, channel_size);
std::memcpy(input_host + channel_size/sizeof(float), channels[1].data, channel_size);
std::memcpy(input_host + 2*channel_size/sizeof(float), channels[2].data, channel_size);

二、输出数据处理（重构为目标形状） #

1. 核心逻辑 #

TensorRT输出缓冲区也是一维线性存储，需根据输出维度的顺序（如(1,32,53,8)对应的N/C/H/W顺序），通过索引计算访问对应位置的数据，或封装成可按多维索引访问的结构。

2. 针对(1,32,53,8)输出维度的实现 #

先通过代码确认输出维度顺序，假设为N → C → H → W，可通过以下两种方式处理：

方法1：嵌套vector模拟多维数组

// 注意：需根据模型输出类型调整（示例用float，你的代码用了int，需匹配）
auto output_dims = context->getBindingDimensions(output_idx);
int N = output_dims.d[0];
int C = output_dims.d[1];
int H = output_dims.d[2];
int W = output_dims.d[3];

// 构建四维vector
std::vector<std::vector<std::vector<std::vector<float>>>> output(
    N, std::vector<std::vector<std::vector<float>>>(
        C, std::vector<std::vector<float>>(
            H, std::vector<float>(W)
        )
    )
);

// 填充数据
for (int n = 0; n < N; ++n) {
    for (int c = 0; c < C; ++c) {
        for (int h = 0; h < H; ++h) {
            for (int w = 0; w < W; ++w) {
                // 一维索引计算：n*C*H*W + c*H*W + h*W + w
                int idx = n * C * H * W + c * H * W + h * W + w;
                output[n][c][h][w] = output_buffer.get()[idx];
            }
        }
    }
}

方法2：自定义索引函数直接访问

无需额外内存，直接通过索引计算访问一维数组：

// 定义索引计算函数（需匹配输出维度顺序）
inline float get_output(const float* buf, int n, int c, int h, int w, int C, int H, int W) {
    return buf[n * C * H * W + c * H * W + h * W + w];
}

// 使用示例：获取第0个batch、第5个通道、第10行、第3列的数据
float value = get_output(output_buffer.get(), 0, 5, 10, 3, 32, 53, 8);

注意事项

• 若输出是NHWC格式，索引计算需调整为n*H*W*C + h*W*C + w*C + c
• 数据类型必须与模型输出一致：若模型输出为float，不能用int类型的output_buffer，否则会导致数据错误

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