caffe - 将 TensorRT 输出转换为 Caffe

Question

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

Answer 1

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

Answer 2

对于那些现在在这个问题上磕磕绊绊的人，我通过使 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);