TensorRT3.0开发手册 （SamplePlugin:Implementing A Custom Layer）

3.8 SamplePlugin:Implementing A Custom Layer

SamplePlugin例程展示了如何在TensorRT中新增一个用户自定义层。例程中实现了如何将MNIST例程中原本直接cuBLAS库实现的最后全卷积层进行替换。

注：当前Plugin层只支持FP32精度。

3.8.1 Key Concepts

本例中解释的关键概念包括：

Ø  如何创建一个用户自定义层

Ø  如何利用NvCaffeParser集成用户自定义层，并导入runtime。

3.8.2 Implementing The Plugin Interface

每个TensorRT步骤操作都需要利用pluginAPI中实现。下面总体介绍了那些方式与每个步骤相关。

Whencreating the network：

在网络构造过程中，如果ITensor::getDimension()在设置层输出大小，或者其他后面的层时候被调用，这些方式会被使用。此外，当builder运行时也会被调用。

Ø  getNbOutputs()

Ø  getOutPutDimensions()

Bythe builder：

Ø  configure()

Ø  getWorkspaceSize()

Atruntime：

Ø  initialize()当engine context构建的时候。

Ø  enqueue()在前向运算时。

Ø  terminate()当销毁engine context时。

在serialization时：

Ø  getSerializationSize()

Ø  serialize()

3.8.3 Defining The Network

定义网络时，TensorRT需要知道层有哪些输出。

注：例程中的维度时没有batchsize的，与ITensor::getDimensions()返回的维度方式相似。例如，对于一个典型的3维度卷积，维度以{C,W,H}格式给出，返回的值也需要{C,W,H}格式。

下列方式了输出层的内容：

int getNbOutputs() const override

{

return 1;

}

getOutputDimensions函数有三个参数：output index，input dimensions与number of inputs。后面两个参数已经在TensorRT内部计算好了，所以你只需要根据输入维度与现有的索引数，计算输出维度即可。

Dims getOutputDimensions(int index, const Dims* inputDims, int

nbInputDims) override

{

assert(index == 0 && nbInputDims == 1 &&inputDims[0].nbDims

== 3);

assert(mNbInputChannels == inputDims[0].d[0] *

inputDims[0].d[1] *

inputDims[0].d[2]);

return DimsCHW(mNbOutputChannels, 1, 1);

}

3.8.4 Layer Configuration

Builder通过调用网络的configure()来根据输入选择算法。再更复杂的情况下，用户可以根据输入维度决定使用哪种卷积算法。

Configure()的方式只有在build时会调用，所有相关的设置都需要在运行时，保存在plugin的参数中，并进行serialized与de-serialized。

3.8.5 Workspace

TensorRT可以在层执行过程中提供临时的存储空间，为了降低内存使用率，内存空间在层间是共享的。TensorRTbuilder调用getWorkspaceSize()来决定空间需求。本例中没有使用workspace。如果需要使用workspace，workspace将在IExecutionContext创建时分配空间，在运行时传到enqueue()中。

同时用户也可以在GPU上申请运行需要的GPU显存，并且利用destructor释放，在回调中返回0。在TensorRT中使用workspace的优势是在运行时可以多个plugin层共享。

3.8.6 Resource Management

在进行auto-tuning或者运行时创建销毁IExecutionContent，builder调用initialize()与terminate()。他们被用来申请与释放层执行过程中需要的资源。此例中，调用cuDNN与cuBLAS以及一些其他cuDNNtensor的操作。

int initialize() override

{

CHECK(cudnnCreate(&mCudnn));

CHECK(cublasCreate(&mCublas));

CHECK(cudnnCreateTensorDescriptor(&mSrcDescriptor));

CHECK(cudnnCreateTensorDescriptor(&mDstDescriptor));

return 0;

}

virtual void terminate() override

{

CHECK(cublasDestroy(mCublas));

CHECK(cudnnDestroy(mCudnn));

}

3.8.7 Runtime Implementation

当执行层实现时，使用enqueue()方式。传到enqueue()中的batchsize是build时最大的batchsize，也可以比这个小。

注：除了batchsize，维度信息没有传到enqueue()。因此，运行时需要其他维度信息，例如，输出的个数，输出的个数，需要作为层数据serialized。

virtual int enqueue(int batchSize, const void*const *inputs,

void**

outputs, void* workspace, cudaStream_t stream)override

{

int nbOutputChannels = mBiasWeights.count;

int nbInputChannels = mKernelWeights.count /

nbOutputChannels;

float kONE = 1.0f, kZERO = 0.0f;

cublasSetStream(mCublas, stream);

cudnnSetStream(mCudnn, stream);

CHECK(cublasSgemm(mCublas, CUBLAS_OP_T, CUBLAS_OP_N,

nbOutputChannels, batchSize, nbInputChannels,&kONE,

reinterpret_cast

float*>(mKernelWeights.values),

nbInputChannels,

reinterpret_cast(inputs[0]),

nbInputChannels, &kZERO,

reinterpret_cast(outputs[0]),

nbOutputChannels));

CHECK(cudnnSetTensor4dDescriptor(mSrcDescriptor,

CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1,nbOutputChannels, 1, 1));

CHECK(cudnnSetTensor4dDescriptor(mDstDescriptor,

CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batchSize,nbOutputChannels,

1, 1));

CHECK(cudnnAddTensor(mCudnn, &kONE,mSrcDescriptor,

mBiasWeights.values, &kONE, mDstDescriptor,outputs[0]));

return 0;

}

3.8.8 Serialization

层参数可以与余下的网络一起进行序列化，序列化系统调用下列函数：

virtual size_t getSerializationSize() override

{

// 3 integers (number of input channels, number ofoutput

channels, bias size), and then the weights:

return sizeof(int)*3 +mKernelWeights.count*sizeof(float) +

mBiasWeights.count*sizeof(float);

}

virtual void serialize(void* buffer) override

{

char* d = reinterpret_cast(buffer), *a =d;

write(d, mNbInputChannels);

write(d, mNbOutputChannels);

write(d, (int)mBiasWeights.count);

serializeFromDevice(d, mKernelWeights);

serializeFromDevice(d, mBiasWeights);

assert(d == a + getSerializationSize());

}

反序列化通过下列构造器实现：

// create the plugin at runtime from a byte stream

FCPlugin(const void* data, size_t length)

{

const char* d = reinterpret_cast(data), *a = d;

mNbInputChannels = read(d);

mNbOutputChannels = read(d);

int biasCount = read(d);

mKernelWeights = deserializeToDevice(d,mNbInputChannels *

mNbOutputChannels);

mBiasWeights = deserializeToDevice(d, biasCount);

assert(d == a + length);

}

3.8.9 Adding The Plugin Into A Network

有三个在网络中新增plugin方式：

1、  定义网络时使用INetwork::addPlugin()函数。

2、  使用parser创建网络。

3、  在build后反序列化网络。

AddPlugin()使用方式见TensorRT API。

3.8.9.1 通过NvCaffeParser创建Plugins

使用NvCaffeParser新增用户自定义网络时，通过nvcaffeParser::IPluginFactory结构创建一个工厂，将实例传到ICaffeParser::parse()中。CreatePlugin()接受层名称，NVCaffe模型文件中提取的一系列权重，这些参数通过层constructor传入。名称是用来区分多个插件的依据，目前除了网络权重不能提取其他NVCaffe网络描述，因此，这些参数都必须在factory中指定。

bool isPlugin(const char* name) override

{

return !strcmp(name, "ip2");

}

virtual nvinfer1::IPlugin* createPlugin(const char*layerName,

const

nvinfer1::Weights* weights, int nbWeights) override

{

// there's no way to pass parameters through from themodel

definition, so we have to define it here explicitly

static const int NB_OUTPUT_CHANNELS = 10;

assert(isPlugin(layerName) && nbWeights == 2&&

weights[0].type ==

DataType::kFLOAT && weights[1].type ==DataType::kFLOAT);

assert(mPlugin.get() == nullptr);

mPlugin = std::unique_ptr(newFCPlugin(weights,

nbWeights, NB_OUTPUT_CHANNELS));

return mPlugin.get();

}

3.8.9.2 Creating Plugins At Runtime

为了在运行时整合用户自定义网络，使用nvinfer1::IPlugin接口，将实例的结构传到IInferRuntime::deserializeCudaEngine()。

// deserialization plugin implementation

IPlugin* createPlugin(const char* layerName, constvoid*

serialData,

size_t serialLength) override

{

assert(isPlugin(layerName));

assert(mPlugin.get() == nullptr);

mPlugin = std::make_unique(serialData,

serialLength);

return mPlugin.get();

}

当使用NvCaffeParser构建，运行时反序列化，层实现假设数据通过NVCaffeParser的权重或者运行时的二进制流传送直至调用initialize()函数，允许数据复制到GPU中。