将TensorRT输出转换为Caffe

Question

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

Answer 1

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

Answer 2

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