2

My network contains some specific layers which are not supported by current tensorRT. so I want to run the conv layers and pooling layers on tensorRT and then use the output from tensorRT as the input of my caffe model which contains some specific layers. Is there some API or example code thar I can refer to? Thanks

Robert Crovella
  • 143,785
  • 11
  • 213
  • 257
wanglinski
  • 114
  • 9

2 Answers2

0

See the source code in the samples directory of your TensorRT installation.

Ashwin Nanjappa
  • 76,204
  • 83
  • 211
  • 292
0

For those stumbling now on this issue I got this to work by making the input and output of TensorRT inference the mutable_gpu_data from caffe blobs:

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);
florentbuisson
  • 526
  • 5
  • 12