Hi all,
I have written a very simple custom plugin in TensorRT which uses an Nvidia Performance Primitive (NPP) function to upsample tensors with stride=2. I’m relatively new to TensorRT and was wondering if using NPP within TensorRT custom plugins is a problem in any way, as none of the examples provided used NPP functions. However, I’ve tested the layer and proven that it works as expected.
Here is the enqueue function (method within the plugin class)
int enqueue(int batchSize, const void* const* inputs, void** outputs, void*, cudaStream_t stream) override
{
size_t inputSizeBytes = this->mCHW.d[0] * this->mCHW.d[1] * this->mCHW.d[2] * sizeof(float);
size_t inputImageSize = this->mCHW.d[1] * this->mCHW.d[2];
size_t outputImageSize = 4 * this->mCHW.d[1] * this->mCHW.d[2]; // 4 = stride^2
cudaError_t ret;
const float* input = reinterpret_cast<const float*>(inputs[0]);
float* output = reinterpret_cast<float*>(outputs[0]);
//input and output definitions for the NPP resize function
NppiSize oSrcSize;
oSrcSize.width = this->mCHW.d[1];
oSrcSize.height = this->mCHW.d[2];
NppiRect oSrcRectROI;
oSrcRectROI.x = 0;
oSrcRectROI.y = 0;
oSrcRectROI.width = this->mCHW.d[1];
oSrcRectROI.height = this->mCHW.d[2];
int nSrcStep = sizeof(float) * this->mCHW.d[1];
NppiSize oDstSize;
oDstSize.width = 2 * this->mCHW.d[1];
oDstSize.height = 2 * this->mCHW.d[2];
NppiRect oDstRectROI;
oDstRectROI.x = 0;
oDstRectROI.y = 0;
oDstRectROI.width = 2 * this->mCHW.d[1];
oDstRectROI.height = 2 * this->mCHW.d[2];
int nDstStep = sizeof(float) * 2 * this->mCHW.d[1];
int eInterpolation = NPPI_INTER_NN ; //Nearest neighbor interpolation
//Iterate through all the feature maps
int ch;
for (ch = 0; ch < this->mCHW.d[0]; ch++)
{
Update the position in the buffer
const Npp32f* pSrc = reinterpret_cast<const Npp32f*>(input + ch*inputImageSize);
Npp32f* pDst = reinterpret_cast<Npp32f*>(output + ch*outputImageSize);
NppStatus status = nppiResize_32f_C1R(pSrc, nSrcStep, oSrcSize, oSrcRectROI,
pDst, nDstStep, oDstSize, oDstRectROI,
eInterpolation);
}
return 0;
}