TensorRT - Custom Layer

2019 May 05 tensorrt

TensorRT - Custom Layer

Custom Layer 的实现

Custom Layer 通过继承 IPlugin类，对 TensorRT 的五个阶段分别实现。

• 创建网络阶段
• 构建 Engine 阶段
• 序列化 Engine 阶段
• 反序列化 Engine 阶段
• 执行推理 Engine 阶段

创建网络阶段

网络创建需要告诉TensorRT该Custom Layer的output矩阵的维数，主要通过以下两种方法传递：

• getNbOutputs()——需要返回该层output矩阵的总数量
• getOutputDimensions()——需要提供该层指定output矩阵的维数

其中该函数有三个参数，分别是output矩阵索引、input矩阵大小以及输入矩阵数。后两者参数是由TensorRT提供的，该方法内部需要计算输出矩阵维数。

构建 Engine 阶段

• configure()——如果需要对该层的input矩阵进行验证或选择卷积算法等等操作可以通过配置此函数实现，如果无需配置任何算法可以留空。

  解释：构建Builder通过调用configure()函数来选择算法，如 input 矩阵验证算法、卷积算法等等。configure() 方法只在构建阶段会被调用，因此任何在该函数中确定的参数需要被存储作为 IPlugin 类的成员变量以供后面使用。

• getWorkspaceSize()——如果在运行时需要TensorRT管理的workspace空间可以在此配置，无需申请返回0即可。

  解释：TensorRT 可以提供 workspace 作为每层网络执行时的临时存储空间，该空间是共享的以减少显存占用。构建 Builder 通过调用 getWorkspaceSize() 来确定 workspace 的需求。如果在该函数中设置了 workspace 则在创建执行 Context 的时候分配该 workspace，在 runtime 是提供给 enqueue() 方法使用，并在销毁时回收。使用 TensorRT 分配的 workspace 的好处在于可以在不同 Custom Layer 之间共享使用以减少显存开销。

序列化 Engine 阶段

• getSerializationSize()——返回序列化所需的缓存空间。

• serialize()——实现该层参数的序列化，参考代码如下：

virtual void serialize(void* buffer) override
{
    char* d = reinterpret_cast<char*>(buffer), *a=d;

    write(d, mNbInputChannels);
    write(d, mNbOutputChannels);
    write(d, (int)mBiasWeights.count);
    serializeFromDevice(d, mKernelWeights);
    serializeFromDevice(d, mBiasWeights);

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

void serializeFromDevice(char*& hostBuffer, Weights deviceWeights)
{
    cudaMemcpy(hostBuffer, deviceWeights.values,
            deviceWeights.count * sizeof(float), cudaMemcpyDeviceToHost);
    hostBuffer += deviceWeights.count * sizeof(float);
}

反序列化 Engine 阶段

Custom Layer 的反序列化过程是在反序列化整个模型时候完成的，并通过实例化该 Plugin 类来实现的，参考代码如下所示：

// deserialize the engine
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(gieModelStream->data(),
                                    gieModelStream->size(), &pluginFactory);

FCPlugin(const void* data, size_t length)
{
    const char* d = reinterpret_cast<const char*>(data), *a =d;
    mNbInputChannels = read<int>(d);
    mNbOutputChannels = read<int>(d);
    int biasCount = read<int>(d);

    mKernelWeights = deserializeToDevice(d, 
                            mNbInputChannels * mNbOutputChannels);
    mBiasWeights = deserializeToDevice(d, biasCount);
    assert(d == a + length);
}

其中的 deserilizationToDevice 函数是该类的私有函数，也需要编写。

执行推理 Engine 阶段

• initialize() and terminate()——分别会在构建 builder 执行网络自动优化和创建/销毁执行 Context 时被调用。

  解释：更多情况下作为其他的资源管理使用，例如 cudnn 资源的初始化和销毁使用

• enqueue()——实现的是该Custom Layer在执行时具体的计算操作。

  解释：该函数拥有五个参数，分别是 input batchsize、input tensor、output tensor、workspace、stream。其中 input batchsize 最大为 Engine 构建时设置的最大 Batchsize，workspace 即前面构建阶段 getWorkspaceSize() 申请的 workspace。需要注意的是这边只传递了 batchsize 信息，但输入以及输出矩阵的维数信息并没有进行传递，因此需要在序列化时将所需的信息进行序列化，例如输入输出矩阵的 channel 数。

TensorRT 感知 Custom Layer

TensorRT 需要在导入网络模型至 TensorRT 和在创建 Runtime Engine（反序列化 Engine）的时候感知到新的Custom Layer。

模型导入时

C++ API

使用 C++ 定义网络结构时使用 addPlugin() 函数即可引入 Custom Layer，该函数的输入为该层的输入矩阵，输入矩阵总数以及该层类实现。

Python API

使用 Python 定义网络结构时使用 add_plugin() 函数即可引入Custom Layer，该函数的输入为该层的输入矩阵，输入矩阵总数以及该层类实现。

NvCaffeParser

NvCaffeParser 的导入方式采用的半自动的导入方式（提供 caffemodel 和 prototxt 的形式），因此无法使用 addPlugin() 函数来显性定义。Caffe的网络层定义中分别有 Name 和 Type 两种参数，在此处为了定位新增的 Custom Layer，TensorRT 中将使用 Name 参数作为标识符。此处的实现是通过继承 nvcaffeparser1::IPluginFactory 类并实现函数 isPlugin() 和 nvinfer1::IPlugin* createPlugin() 函数并将实例化的新类传递给 ICaffeParser::parse() 来完成的。

• isPlugin()——判断是否是Custom Layer。

  解释：该函数输入为 CaffeParser 传递过来的网络层Name参数，因此通过 Name 参数比对即可判断是否该层为特殊的 Custom Layer。

• nvinfer1::IPlugin* createPlugin()——定义层参数。

  解释：目前暂时无法通过模型定义传递参数信息，因此需要在此层中显性定义，其中该层的输入为网络层 Name 和模型 Weights。

需要注意的是，采用 NvCaffeParser 导入的方式时，Custom Layer 的实现将默认使用从 NvCaffeParser 传递的 Weights 并认为在 initialize() 函数执行前数据已经准备就绪。参考实现如下：

// caffe parser plugin implementation
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 the model definition, 
    // so we have to define it here exlicitly
    static const int NB_OUT_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<FCPlugin>(new FCPlugin(weights, nbWeights, 
                                        NB_OUTPUT_CHANNELS));
    return mPlugin.get();
}

ICaffeParser* parser = createCaffeParser();
parser->setPluginFactory(pluginFactory);

创建 Runtime Engine 时

创建 Runtime Engine 时（反序列化Engine时）和模型导入时不一样，此处无论是哪种方式导入的模型都已经转化成优化好的序列化的 TensorRT Engine。因此此时只需要继承类 nvinfer1::IPluginFactory 并实现其接口函数 nvinfer1::IPlugin，之后将继承的新类实例化并传递给 IInferRuntime::deserializeCudaEngine() 即可，参考代码如下：

// deserialization plugin implementation
IPlugin* createPlugin(const char* layerName, const void* serialData, 
                      size_t serialLength) override
{
    assert(isPlugin(layerName));
    assert(mPlugin.get() == nullptr);
    mPlugin = std::unique_ptr<FCPlugin>(new FCPlugin(serialData, 
                                        serialLength));
    return mPlugin.get();
}

// deserialize the engine
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(gieModelStream->data(),
                                    gieModelStream->size(), &pluginFactory);

需要注意的是，在创建 Runtime时（反序列化的实现如前文 “Custom Layer实现” 中反序列化部分所示），这部分的实现默认使用字节流数据并认为认为在 initialize() 函数执行前数据已经准备就绪。