将TensorRT输出转换为Caffe

Question

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

Answer 1

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

Answer 2

对于那些在这个问题上磕磕绊绊的人，我通过将来自caffe blobs的TensorRT推理的输入和输出推断为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);