tensorrt 自定义层详解

公用函数

宏定义

static const char* G_PRELU_TYPE = "PReLU";
static const char* G_PRELU_NAME = "PReLU_TRT";

宏函数

#define CUDA_KERNEL_LOOP(i, n) \for (int i = blockIdx.x * blockDim.x + threadIdx.x; \i < (n); \i += blockDim.x * gridDim.x)

核相关定义

static const int CUDA_NUM_THREADS = 512;// CUDA: number of blocks for threads.
inline int CAFFE_GET_BLOCKS(const int N) {
    return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}

class PReLUPlugin详解

私有成员

输入的CHW,值,类型及GPU端拷贝

int mNbInputChannels, mNbInputHeight, mNbInputWidth;
nvinfer1::Weights mWeights;
nvinfer1::DataType mDataType{
    nvinfer1::DataType::kFLOAT};
void* mDeviceKernel{
    nullptr};

PReLUPlugin(const nvinfer1::Weights* weights, int nbWeight);

这个构造函数的作用是将此层的参数(权重,偏置等参数)读取到内部
将模型文件参数的值复制给私有成员

PReLUPlugin::PReLUPlugin(const nvinfer1::Weights *weights, int nbWeights) {
    mWeights = weights[0];mWeights.values = malloc(mWeights.count * type2size(mWeights.type));memcpy(const_cast<void *>(mWeights.values), weights[0].values, mWeights.count * type2size(mWeights.type));
}

PReLUPlugin::PReLUPlugin(const void *data, size_t length)

从序列化模型文件里读取数据
将参数的CHW,nvinfer1::Weights结构体内的变量copy到私有成员
最后assert内存大小是否相同的

PReLUPlugin::PReLUPlugin(const void *data, size_t length) {
    const char *d = static_cast<const char *>(data), *a = d;read<int>(d, mNbInputChannels);read<int>(d, mNbInputHeight);read<int>(d, mNbInputWidth);read<nvinfer1::DataType>(d, mDataType);read<int64_t>(d, mWeights.count);read<nvinfer1::DataType>(d, mWeights.type);mWeights.values = nullptr;mWeights.values = malloc(mWeights.count * type2size(mWeights.type));memcpy(const_cast<void *>(mWeights.values), d, mWeights.count * type2size(mWeights.type));d = d + mWeights.count * type2size(mWeights.type);ASSERT(d == a + length);

virtual size_t getSerializationSize() const override

返回下在buffer里占用的大小,其实就是读取序列化构造函数里面所有变量的内存大小

size_t PReLUPlugin::getSerializationSize() const {
    return sizeof(mNbInputChannels) + sizeof(mNbInputWidth) + sizeof(mNbInputHeight) + sizeof(mDataType) + sizeof(mWeights.count) + sizeof(mWeights.type) + mWeights.count * type2size(mWeights.type);
}

virtual void serialize(void* buffer) const override;

序列化插件到buffer 给构造函数相反一个读一个写,读写的参数相同

void PReLUPlugin::serialize(void *buffer) const {
    char *d = static_cast<char *>(buffer), *a = d;write(d, mNbInputChannels);write(d, mNbInputHeight);write(d, mNbInputWidth);write(d, mDataType);write(d, mWeights.count);write(d, mWeights.type);convertAndCopyToBuffer(d, mWeights, mWeights.type);ASSERT(d == a + getSerializationSize());
}

PReLUPlugin() = delete;

无参数构造函数无意义

~PReLUPlugin();

析构函数
释放你申请的空间防止段错误

PReLUPlugin::~PReLUPlugin() {
    if (mWeights.values) {
    free(const_cast<void *>(mWeights.values));mWeights.values = nullptr;}if (mDeviceKernel) {
    cudaFree(mDeviceKernel);mDeviceKernel = nullptr;}
}

virtual int getNbOutputs() const override;

返回输出tensor的数量, 比如说prelu,输出个数跟relu一样是1,这个取决于你的自定义层,大概是输入一个变成几个,待研究,后续补充

int PReLUPlugin::getNbOutputs() const {
    return 1;
}

virtual nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) override

返回输出tensor的维度,根据输入维度和操作自行计算
index这边有疑问?待解决

nvinfer1::Dims PReLUPlugin::getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) {
    if(index == 0) {
    return nvinfer1::Dims3(inputs[0].d[0],inputs[0].d[1],inputs[0].d[2]);} // else if(index == n) {
    // for other outputs if exists.// }else {
    ASSERT(false);}
}

virtual bool supportsFormat(const nvinfer1::DataType type, nvinfer1::PluginFormat format) const override;

看看你支持什么格式,half,int

bool PReLUPlugin::supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format) const {
    return (type == nvinfer1::DataType::kFLOAT | type == nvinfer1::DataType::kHALF) && format == nvinfer1::PluginFormat::kNCHW;
}

virtual void configureWithFormat(const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims,int nbOutputs, nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) override;

判断下支持的格式,并给私有成员赋值

void PReLUPlugin::configureWithFormat(const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims, int nbOutputs,nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) {
    ASSERT((type == nvinfer1::DataType::kFLOAT | type == nvinfer1::DataType::kHALF)&& format == nvinfer1::PluginFormat::kNCHW);mNbInputChannels = inputDims[0].d[0]; mNbInputHeight = inputDims[0].d[1];mNbInputWidth = inputDims[0].d[2];mDataType = type;
}

virtual int initialize() override

初始化你的插件,其实就是初始化好gpu context, 将你的权重从内存拷贝到gpu上,如果设定了fp16,当然也要先做转化在拷贝到gpu

int PReLUPlugin::initialize() {
    convertAndCopyToDeivce(mDeviceKernel, mWeights, mDataType);return 0;
}

virtual void terminate() override

释放内存和显存,给析构函数一样,只不过调用对象不同

void PReLUPlugin::terminate() {
    if (mWeights.values){
    free(const_cast<void *>(mWeights.values));mWeights.values = nullptr;}if (mDeviceKernel){
    cudaFree(mDeviceKernel);mDeviceKernel = nullptr;}
}

virtual size_t getWorkspaceSize(int maxBatchSize) const override

很难解释, 直接返回0即可

size_t PReLUPlugin::getWorkspaceSize(int maxBatchSize) const
{
    return 0;
}

virtual const char* getPluginType() const override

tensorrt内部定义的

const char *PReLUPlugin::getPluginType() const {
    return G_PRELU_TYPE;
}

virtual const char* getPluginVersion() const override;

tensorrt内部定义的

const char *PReLUPlugin::getPluginVersion() const {
    return G_PLUGIN_VERSION;
}

virtual void destroy()

tensorrt内部接口来析构

void PReLUPlugin::destroy() {
    delete this; 
}

virtual nvinfer1::IPluginV2* clone() const override;

nvinfer1::IPluginV2* PReLUPlugin::clone() const {
    return new PReLUPlugin(&mWeights, 1);
}

virtual void setPluginNamespace(const char* pluginNamespace) override {}

不用实现留空

virtual const char* getPluginNamespace() const override

const char* PReLUPlugin::getPluginNamespace() const {
    return G_PLUGIN_NAMESPACE;
}

virtual int enqueue(int batchSize, const voidconst * inputs, void** outputs, void workspace, cudaStream_t stream) override

此层功能的代码实现,此处调用推理函数

int PReLUPlugin::enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream)
{
    const int count = batchSize * mNbInputChannels * mNbInputWidth * mNbInputHeight;const int channels = mNbInputChannels;const int dim = mNbInputWidth * mNbInputHeight;const int div_factor = 1;if (mDataType == nvinfer1::DataType::kFLOAT){
    const float zerof{
    0.0f};CUDA_CHECK(Forward_gpu(count, channels, dim,reinterpret_cast<const float *>(mDeviceKernel),reinterpret_cast<const float *>(inputs[0]),reinterpret_cast<float *>(outputs[0]),zerof,div_factor,stream));} else {
    const __half zeroh = __half(0.0f);CUDA_CHECK(Forward_gpu(count, channels, dim,reinterpret_cast<const __half *>(mDeviceKernel),reinterpret_cast<const __half *>(inputs[0]),reinterpret_cast<__half *>(outputs[0]),zeroh,div_factor,stream));}return 0;
}

核函数调用

template <typename Ftype>
cudaError_t Forward_gpu(const int count, const int channels, const int dim,const Ftype* mDeviceKernel,const Ftype* bottom_data, Ftype* top_data, const Ftype zero,const int div_factor, const cudaStream_t stream) {
    PReLUForward<<<CAFFE_GET_BLOCKS(count), CUDA_NUM_THREADS, 0, stream>>>(count, channels, dim, mDeviceKernel, bottom_data, top_data, zero, div_factor);cudaError_t err = cudaGetLastError();return err;
}

template <typename Ftype>
__global__ void PReLUForward(const int n, const int channels, const int dim,const Ftype* slope_data,const Ftype* in, Ftype* out,const Ftype zero,const int div_factor) {
    CUDA_KERNEL_LOOP(index, n) {
    int c = (index / dim) % channels / div_factor;if(in[index] > zero) {
    out[index] = in[index];} else {
    out[index] = in[index] * slope_data[c];}}
}