TensorRT学习（实战-自定义算子）

YOLOv4进行TensorRT推理的时候会使用Mish激活函数，而使用到的mish激活函数没有在TensorRT进行实现。故需要进行实现对应TensorRT插件，故需要进行Mish激活函数的实现。

Mish激活函数定义

def mish_fun(x):
    tmp = np.log(1 + np.exp(x))
    tmp = np.tanh(tmp)
    tmp = tmp * x
    return tmp

如上所示激活函数的表达式，该激活还是比较复杂的，需要实现对应的tensorRT插件

tensorRT 插件

使用C++ API添加自定义图层
您可以通过从TensorRT的插件基类之一派生来实现自定义层。
从插件的一个基类派生插件类。它们在支持具有不同类型/格式的I/O或具有动态形状的网络方面具有不同的表达能力。下表总结了基类，按表达性从低到高的顺序排列。
注意：如果插件是用于一般用途，请提供FP32实现，以便允许它在任何网络上正常运行。

为了在网络中使用插件，您必须首先在TensorRT的PluginRegistry（C++，Python）中注册它。不是直接注册插件，而是注册插件的工厂类的实例，从PluginCreator（C++，Python）派生。plugin creator类还提供了有关插件的其他信息：其名称、版本和插件字段参数。

有两种方法可以在注册表中注册插件：

TensorRT提供了一个宏REGISTER_TENSORT_PLUGIN，用于在注册表中静态注册插件创建者。请注意，NREGISTER_TENSORT_PLUGI始终在默认名称空间（“”）下注册创建者。
通过创建您自己的入口点（类似于initLibNvInferPlugins）并在插件注册表上调用registerCreator来registerCreator。这比静态注册更好，因为它提供了潜在的更低的内存占用，并允许插件在唯一的名称空间下注册。这确保了在不同插件库之间的构建时间期间没有名称冲突。
调用IPluginCreator::createPlugin（）返回一个IPluginV2类型的插件对象。您可以使用addPluginV2（）将插件添加到TensorRT网络，这将使用给定的插件创建网络层。

例如，您可以向网络添加插件层，如下所示：

// Look up the plugin in the registry
auto creator = getPluginRegistry()->getPluginCreator(pluginName, pluginVersion);
const PluginFieldCollection* pluginFC = creator->getFieldNames();
// Populate the fields parameters for the plugin layer 
// PluginFieldCollection *pluginData = parseAndFillFields(pluginFC, layerFields); 
// Create the plugin object using the layerName and the plugin meta data
IPluginV2 *pluginObj = creator->createPlugin(layerName, pluginData);
// Add the plugin to the TensorRT network 
auto layer = network.addPluginV2(&inputs[0], int(inputs.size()), pluginObj);
… (build rest of the network and serialize engine)
// Destroy the plugin object
pluginObj->destroy()
… (free allocated pluginData)

注意：前面描述的createPlugin方法在堆上创建了一个新的插件对象，并返回一个指向它的指针。确保销毁pluginObj，如前所示，以避免内存泄漏。
在序列化期间，TensorRT引擎在内部存储所有IPluginV2类型插件的插件类型、插件版本和命名空间（如果存在）。在反序列化过程中，TensorRT从插件注册表中查找插件创建者，并调用IPluginCreator::deserializePlugin（）。当引擎被删除时，引擎通过调用IPluginV2::destroy()方法销毁在引擎构建期间创建的插件对象的克隆。您有责任确保您创建的插件对象在添加到网络后被释

注：

不要序列化所有插件参数：只有那些插件在运行时正确运行所需的。可以省略生成时间参数。
以相同的顺序序列化和反序列化插件参数。在反序列化过程中，验证插件参数是否初始化为默认值或反序列化值。未初始化的参数会导致未定义的行为。
如果您是汽车安全用户，则必须调用getSafePluginRegistry（）而不是getPluginRegistry（）。还必须使用REGISTER_SAFE_TENSORT_PLUGIN宏，而不是REGISTER_TENSORT_PLUGIN。

关键类说明

IPluginV2Ext

用户实现层的插件创建类。

IPluginCreator

用户实现层的插件创建器类。

Mish激活函数

#ifndef _MISH_PLUGIN_H
#define _MISH_PLUGIN_H

#include <string>
#include <vector>
#include "NvInfer.h"

namespace nvinfer1
{
    class MishPlugin: public IPluginV2IOExt
    {
        public:
            // 显示的构造函数
            explicit MishPlugin();
            // 构造函数
            MishPlugin(const void* data, size_t length);
            ~MishPlugin();
            // 返回plugin的输出数量
            int getNbOutputs() const override
            {
                return 1;
            }
            // 输出张量的输出的维度
            Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;
            // 初始化要执行的层。这在引擎创建时调用
            int initialize() override;
            // 释放插件层初始化过程中获取的资源
            virtual void terminate() override {};
            // 获得该层工作空间的大小
            virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;}
            // 执行该的层的处理
            virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override;
            //  查找所需的序列化缓冲区的大小
            virtual size_t getSerializationSize() const override;
            // 进行对应的序列化
            virtual void serialize(void* buffer) const override;

            bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override {
                return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
            }

            const char* getPluginType() const override;

            const char* getPluginVersion() const override;
            //   销毁插件对象。这将在网络、构建器或引擎被销毁时调用
           void destroy() override;
            // 克隆对象
            IPluginV2IOExt* clone() const override;
            // 这是对象所属的命名空间
            void setPluginNamespace(const char* pluginNamespace) override;
            // 返回对象的命名空间
            const char* getPluginNamespace() const override;
            // 返回对象输出时间
            DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;
            // 是否进行舒畅张量在批处理中广播
            bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
            // 如果插件可以使用跨批广播的输入而无需刻意的复制
            bool canBroadcastInputAcrossBatch(int inputIndex) const override;
            // 将插件对象附加到执行上下文，
            void attachToContext(
                    cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;
            // 使用输入和输出数据类型配置网络层
            void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;
            // 从执行上下文中分离插件对象.
            void detachFromContext() override;

            int input_size_;
        private:
            void forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize = 1);
            int thread_count_ = 256;
            const char* mPluginNamespace;
    };

    class MishPluginCreator : public IPluginCreator
    {
        public:
            // 构造函数
            MishPluginCreator();
            // 析构函数
            ~MishPluginCreator() override = default;
            // 获得插件的名字
            const char* getPluginName() const override;
            // 获得插件的版本
            const char* getPluginVersion() const override;
            // 返回需要传递给createPlugin的字段列表
            const PluginFieldCollection* getFieldNames() override;
            // 返回一个对应的插件对象
            IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
            // 在插件层的反序列化过程中调用。返回插件对象
            IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
            // 根据插件所属的插件库设置插件创建者的命名空间
            void setPluginNamespace(const char* libNamespace) override
            {
                mNamespace = libNamespace;
            }
            // 返回插件创建者对象的命名空间
            const char* getPluginNamespace() const override
            {
                return mNamespace.c_str();
            }

        private:
            std::string mNamespace;
            static PluginFieldCollection mFC;
            static std::vector<PluginField> mPluginAttributes;
    };
    REGISTER_TENSORRT_PLUGIN(MishPluginCreator);
};
#endif 

#include <cmath>
#include <stdio.h>
#include <cassert>
#include <iostream>
#include "mish.h"

namespace nvinfer1
{
    MishPlugin::MishPlugin()
    {
    }

    MishPlugin::~MishPlugin()
    {
    }

    // create the plugin at runtime from a byte stream
    MishPlugin::MishPlugin(const void* data, size_t length)
    {
        assert(length == sizeof(input_size_));
        input_size_ = *reinterpret_cast<const int*>(data);
    }

    void MishPlugin::serialize(void* buffer) const
    {
        *reinterpret_cast<int*>(buffer) = input_size_;
    }

    size_t MishPlugin::getSerializationSize() const
    {  
        return sizeof(input_size_);
    }

    int MishPlugin::initialize()
    { 
        return 0;
    }

    Dims MishPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
    {
        assert(nbInputDims == 1);
        assert(index == 0);
        input_size_ = inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2];
        // Output dimensions
        return Dims3(inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]);
    }

    // Set plugin namespace
    void MishPlugin::setPluginNamespace(const char* pluginNamespace)
    {
        mPluginNamespace = pluginNamespace;
    }

    const char* MishPlugin::getPluginNamespace() const
    {
        return mPluginNamespace;
    }

    // Return the DataType of the plugin output at the requested index
    DataType MishPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
    {
        return DataType::kFLOAT;
    }

    // Return true if output tensor is broadcast across a batch.
    bool MishPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
    {
        return false;
    }

    // Return true if plugin can use input that is broadcast across batch without replication.
    bool MishPlugin::canBroadcastInputAcrossBatch(int inputIndex) const
    {
        return false;
    }

    void MishPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput)
    {
    }

    // Attach the plugin object to an execution context and grant the plugin the access to some context resource.
    void MishPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)
    {
    }

    // Detach the plugin object from its execution context.
    void MishPlugin::detachFromContext() {}

    const char* MishPlugin::getPluginType() const
    {
        return "Mish_TRT";
    }

    const char* MishPlugin::getPluginVersion() const
    {
        return "1";
    }

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

    // Clone the plugin
    IPluginV2IOExt* MishPlugin::clone() const
    {
        MishPlugin *p = new MishPlugin();
        p->input_size_ = input_size_;
        p->setPluginNamespace(mPluginNamespace);
        return p;
    }

    __device__ float tanh_activate_kernel(float x){return (2/(1 + expf(-2*x)) - 1);}

    __device__ float softplus_kernel(float x, float threshold = 20) {
        if (x > threshold) return x;                // too large
        else if (x < -threshold) return expf(x);    // too small
        return logf(expf(x) + 1);
    }

    __global__ void mish_kernel(const float *input, float *output, int num_elem) {

        int idx = threadIdx.x + blockDim.x * blockIdx.x;
        if (idx >= num_elem) return;

        //float t = exp(input[idx]);
        //if (input[idx] > 20.0) {
        //    t *= t;
        //    output[idx] = (t - 1.0) / (t + 1.0);
        //} else {
        //    float tt = t * t;
        //    output[idx] = (tt + 2.0 * t) / (tt + 2.0 * t + 2.0);
        //}
        //output[idx] *= input[idx];
        output[idx] = input[idx] * tanh_activate_kernel(softplus_kernel(input[idx]));
    }

    void MishPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) {
        int block_size = thread_count_;
        int grid_size = (input_size_ * batchSize + block_size - 1) / block_size;
        mish_kernel<<<grid_size, block_size>>>(inputs[0], output, input_size_ * batchSize);
    }

    int MishPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream)
    {
        //assert(batchSize == 1);
        //GPU
        //CUDA_CHECK(cudaStreamSynchronize(stream));
        forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize);
        return 0;
    }

    PluginFieldCollection MishPluginCreator::mFC{};
    std::vector<PluginField> MishPluginCreator::mPluginAttributes;

    MishPluginCreator::MishPluginCreator()
    {
        mPluginAttributes.clear();

        mFC.nbFields = mPluginAttributes.size();
        mFC.fields = mPluginAttributes.data();
    }

    const char* MishPluginCreator::getPluginName() const
    {
            return "Mish_TRT";
    }

    const char* MishPluginCreator::getPluginVersion() const
    {
            return "1";
    }

    const PluginFieldCollection* MishPluginCreator::getFieldNames()
    {
            return &mFC;
    }

    IPluginV2IOExt* MishPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
    {
        MishPlugin* obj = new MishPlugin();
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }

    IPluginV2IOExt* MishPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
    {
        // This object will be deleted when the network is destroyed, which will
        // call MishPlugin::destroy()
        MishPlugin* obj = new MishPlugin(serialData, serialLength);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }

}

将该对象添加到模型中

Weights emptywts{DataType::kFLOAT, nullptr, 0};
    //卷积层处理
        //!
    //! \brief Add a multi-dimension convolution layer to the network.
    //!
    //! \param  The ininputput tensor to the convolution.
    //! \param nbOutputMaps The number of output feature maps for the convolution.
    //! \param kernelSize The multi-dimensions of the convolution kernel.
    //! \param kernelWeights The kernel weights for the convolution.
    //! \param biasWeights The optional bias weights for the convolution.
    // IConvolutionLayer* addConvolutionNd(
    //     ITensor& input, int32_t nbOutputMaps, Dims kernelSize, Weights kernelWeights, Weights biasWeights) noexcept
    // {
    //     return mImpl->addConvolutionNd(input, nbOutputMaps, kernelSize, kernelWeights, biasWeights);
    // }
    IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap["module_list." + std::to_string(linx) + ".Conv2d.weight"], emptywts);
    assert(conv1);
    // 设置对应参数
    conv1->setStrideNd(DimsHW{s, s});
    conv1->setPaddingNd(DimsHW{p, p});

    // 设置对应的批量归一化数据
    IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "module_list." + std::to_string(linx) + ".BatchNorm2d", 1e-4);

    // 创建对应的mish激活函数
    auto creator = getPluginRegistry()->getPluginCreator("Mish_TRT", "1");
    const PluginFieldCollection* pluginData = creator->getFieldNames();
    IPluginV2 *pluginObj = creator->createPlugin(("mish" + std::to_string(linx)).c_str(), pluginData);
    ITensor* inputTensors[] = {bn1->getOutput(0)};
    auto mish = network->addPluginV2(&inputTensors[0], 1, *pluginObj);
    return mish;