tensorRT-添加自定义层

发表于 | 分类于

TensorRT是一个基于GPU高性能的前向计算的C++库。TensorRT导入网络定义，通过合并tensors与layers，权值量化，来进行网络优化。支持多个深度学习框架，当前，tensorRT3 直接支持的类型的网络层有：

可以说通用的、常见的网络层都已经获得TensorRT中原生支持了，但是在实际应用中可能会采用一些特殊或较新的网络层次，这部分的内容在TensorRT中并未得到原生支持，但对于自定义层又该如何实现，如何载入到tensorRT中去呢。

自定义层实现

tensorRT中提供了一些自定义层的api，根据这个可以定义自己的网络层，，并让tensorRT通过API调用该层，从而像一个插件（Iplugin）一样将新的层添加进TensorRT中。从前一篇博文中可以了解到，tensorRT大致分为5个阶段：创建网络、构建推理Engine、序列化引擎、反序列化引擎、以及执行推理引擎doInference。因此，对于自定义层，我们也需要 针对这5个阶段分别有相应的实现。还好tensorRT提供了一个示例供我们借鉴。在示例samplePlugin.cpp中。查看整个实现，总的来说就是继承IPlugin类来实现的。参考代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#include "NvInfer.h"
using namespace nvinfer1;
class MyPlugin : IPlugin
{
public:
   int getNbOutputs() const;
   Dims getOutputDimensions(int index, const Dims* inputs, 
                            int nbInputDims);
   void configure(const Dims* inputDims, int nbInputs, 
                  const Dims* outputDims, int nbOutputs, 
                  int maxBatchSize);
   int initialize();
   void terminate();
   size_t getWorkspaceSize(int maxBatchSize) const;
   int enqueue(int batchSize, const void* inputs, 
               void** outputs, void* workspace, 
               cudaStream_t stream);
   size_t getSerializationSize();
   void serialize(void* buffer);
protected:
   virtual ~MyPlugin() {}
};

对于上面各个函数实现目的，以及含义，可以参考这里，实现好了自定义层之后，如何让tensorRT感知到这层呢。

TensorRT的感知

请问：对于TensorRT的哪些部分需要感知到新的Custom Layer？
回答：分别是在导入网络模型至TensorRT以及在创建Runtime Engine（反序列化Engine）的时候。这个也可以参考这篇文章的讲解。

最后

自定义层的实现也就上面那么多内容，不过要自己实现有点困难，不过，tensorRT中给出了几个例子，samplePlugin.cpp文件中介绍了实现fc层，以及sampleFasterRCNN.cpp文件也实现了一些自定义层，根据这些例子，实现自定义层，困难应该不大。此外，自定义层暂时只支持FP32精度的，不支持INT8/FP16的。

Python API

这里顺便推荐一下tensorRT的Python api，想了解可以看这里：tensorRT python API