將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);