TensorRT推理输出不完整问题求助

核心问题分析 #

你的代码存在两个关键错误，导致输出后两位为0：

1. 错误使用带对齐的内存分配/拷贝接口
   TensorRT的输入输出张量在GPU内存中是线性连续存储的（按行优先展开为一维数组），但你用cudaMallocPitch分配了带有行间距（pitch）的内存，TensorRT不会按照pitch布局写入数据。后续用cudaMemcpy2D读取时，第二行数据会读取到未初始化的内存区域，因此出现0值。

2. 显式Batch模式下调用了错误的推理接口
   你启用了kEXPLICIT_BATCH模式，但仍使用针对隐式Batch设计的enqueue方法，该方法的batch size参数会被忽略，且无法正确同步流操作。

修改方案 #

1. 替换内存分配/拷贝方式 #

将cudaMallocPitch/cudaMemcpy2D改为普通的cudaMalloc/cudaMemcpy，直接处理连续内存：

// 替换原有的内存分配与拷贝代码
size_t total_size = 2 * 2 * sizeof(float); // 总字节数：4个float
cudaMalloc(&buffers[0], total_size);
cudaMalloc(&buffers[1], total_size);

// 输入拷贝：二维数组在C++中本身是连续内存，可直接拷贝
cudaMemcpy(buffers[0], td_inputTensor, total_size, cudaMemcpyHostToDevice);

// 推理（后续修改enqueue调用）

// 输出拷贝：直接读取连续内存到二维数组
cudaMemcpy(td_outputTensor, buffers[1], total_size, cudaMemcpyDeviceToHost);

2. 改用显式Batch对应的推理接口 #

将enqueue替换为enqueueV2，并使用CUDA流确保操作完成：

// 替换原有的enqueue调用
cudaStream_t stream;
cudaStreamCreate(&stream);
context->enqueueV2(buffers.data(), stream, nullptr);
cudaStreamSynchronize(stream); // 等待推理完成
cudaStreamDestroy(stream);

3. 移除冗余配置（可选） #

在显式Batch模式下，builder->setMaxBatchSize(2)是多余的，可删除该代码，Batch size由context->setBindingDimensions设置的维度控制。

修改后验证 #

运行修改后的代码，输入83 86 77 15时，输出应为-83 -86 -77 -15，符合模型取反的预期。

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