将 TensorRT 输出转换为 Caffe答案

【问题标题】：Convert TensorRT output to Caffe将 TensorRT 输出转换为 Caffe
【发布时间】：2017-04-04 13:07:04
【问题描述】：

我的网络包含一些当前 tensorRT 不支持的特定层。所以我想在 tensorRT 上运行 conv 层和池化层，然后使用 tensorRT 的输出作为我的 caffe 模型的输入，该模型包含一些特定的层。我可以参考一些 API 或示例代码吗？谢谢

【问题讨论】：

标签： caffe nvidia-digits tensorrt

【解决方案1】：

对于那些现在在这个问题上磕磕绊绊的人，我通过让 TensorRT 的输入和输出推断来自 caffe blobs 的 mutable_gpu_data 来实现这一点：

auto* gpuImagePtr = inputBlob->mutable_gpu_data();
cudaMemcpy(gpuImagePtr, inputData, mNetInputMemory, cudaMemcpyHostToDevice);

std::vector<void*> buffers(2);
buffers[0] = gpuImagePtr;
buffers[1] = outputBlob->mutable_gpu_data();

cudaContext->enqueue(batchSize, &buffers[0], stream, nullptr);

【讨论】：

【解决方案2】：

请参阅您的 TensorRT 安装目录的 samples 中的源代码。

【讨论】：