将TensorRT输出转换为咖啡因

Question

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

Answer 1

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

Answer 2

对于那些在这个问题上遇到困难的人，我通过将TensorRT推理的输入和输出设置为来自caffe blob的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);