TensorRT Plugin笔记2

TensorRT Plugin细节

这一周又温习了一遍TensorRT-Plugin,这次重点关注了一遍代码细节,下面是我整合了几篇好的资料做的笔记,分享出来。
TensorRT简单介绍:

1)闭源部分就是官方提供的库,是TRT的核心部分;

2)开源部分在github上,包含Parser(caffe, onnx)、Sample和一些plugin。
在这里插入图片描述
需要写2个类:

1)MyCustomPlugin,继承IPluginV2Ext/IPluginV2IOExt/IPluginV2DynamicExt,是插件类,用于写插件的具体实现;

2)MyCustomPluginCreator,继承BaseCreator,是插件工厂类,用于根据需求创建该插件。

class MyCustomPlugin final : public nvinfer1::IPluginV2DynamicExt
class MyCustomPluginCreator : public BaseCreator

注意:

1)编写plugin,需要继承TRT的base class(下图即base class特性);

在这里插入图片描述

2)Static Shape,用IPluginV2IOExt;Dynamic Shape,则使用IPluginV2DynamicExt。

Static Shape Plugin API

MyCustomPlugin(int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias);  // 构造函数,用于网络定义阶段

MyCustomPlugin(void const* serialData, size_t serialLength);    // 构造函数,用于反序列化阶段

int getNbOutputs() const;      // 获得layer的输出个数

nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims);    // 获得layer的输出维度

nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const;     // 获得输出数据类型

size_t getSerializationSize() const;      //返回序列化时需要写多少字节到buffer中

void serialize(void* buffer) const;      //序列化函数,将plugin的参数权值写入到buffer中

const char* getPluginType() const;       // 获得plugin的type,用于反序列化使用

const char* getPluginVersion() const;       //获得plugin的version,用于反序列化使用

int initialize();          // 初始化函数,在这个插件准备开始run之前执行。一般申请权值显存空间并copy权值

void terminate();         // terminate函数就是释放initialize开辟的一些显存空间

void destroy();           // 释放整个plugin占用的资源

void configurePlugin(const nvinfer1::PluginTensorDesc* in, int nbInput, const nvinfer1::PluginTensorDesc* out, int nbOutput);          // 判断输入是否符合标准

bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const;          // 判断输入、输出的格式

size_t getworkspaceSize(int maxBatchSize) const;        // 获得plugin所需要的显存大小

int enqueue(int batchSize, const void* const* inputs, void** outputs, void* * workspace, cudaStream_t stream);   // 推理函数

const char* setPluginNamespace() const;            // 为这个插件设置namespace名字,每个plugin定义1个专属的Namespace,如果不设置则默认是"",需要注意的是同一个namespace下的plugin如果名字相同会产生冲突
const char* getPluginNamespace() const;   // 获取plugin的命名空间
const PluginFieldCollection *GridAnchorBasePluginCreator::getFieldNames();     // PluginFieldCollection的主要作用是传递插件op所需要的权重和参数
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator);  // 将plugin附加到执行上下文,并授予plugin对某些上下文资源的访问权限
void detachFromContext();      // 将插件对象从其执行上下文中分离出来
构造函数和析构函数
构造函数

构造函数可以写1~3个,通常第一个对应def,第二个对应clone,第三个对应序列化的。

1、用于network definition阶段,PluginCreator创建该插件时调用的构造函数,需要传递权重信息以及参数。也可用于clone阶段,或者再写一个clone构造函数。

MyCustomPlugin(int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias);

​2、clone:顾名思义,就是克隆,将这个plugin对象克隆一份给TensorRT的builder、network或者engine。这个成员函数会调用下面的这个构造函数:

MyCustomPlugin(float in_channel, const std::vector<float>& weight, const std::vector<float>& bias);

将要克隆的plugin的权重和参数传递给这个构造函数。

IPluginV2DynamicExt* MyCustomPlugin::clone() const

{
	auto plugin = new MyCustomPlugin{_in_channel, _weight, _bias};
    plugin->setPluginNamespace(mPluginNamespace);
    return plugin;
}

clone成员函数主要用于传递不变的权重和参数,将plugin复制n多份,从而可以被不同engine、builder、network使用。

3、用于在deserialize阶段,用于将序列化好的权重和参数传入该plugin并创建。

MyCustomPlugin(void const* serialData, size_t serialLength);

注意需要把默认构造函数删掉;

MyCustomPlugin() = delete;
析构函数

析构函数则需要执行terminate,terminate函数就是释放这个op之前开辟的一些显存空间;

MyCustomPlugin::~MyCustomPlugin(){

	terminate();

}
输出相关函数

1、获得layer的输出个数

int getNbOutputs() const;

2、根据输入个数和输入维度,获得第index个输出的维度

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

3、根据输入个数和输入类型,获得第index个输出的类型

nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const;
序列化和反序列化相关函数

1、返回序列化时需要写多少字节到buffer中

size_t MyCustomPlugin::getSerializationSize() const
{
    return (serialized_size(_in_channel) + serialized_size(_weight) + serialized_size(_bias));
};

2、序列化函数,将plugin的参数权值写入到buffer中

void MyCustomPlugin::serialize(void* buffer) const
{
	serialize_value(&buffer, _in_channel);
    serialize_value(&buffer, _weight);
    serialize_value(&buffer, _bias);
};

3、如果这个op使用到了一些其他东西,例如cublas handle,可以直接借助TensorRT内部提供的cublas handle:

void MyCustomPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)

{
	mCublas = cublasContext;
}

4、获得plugin的type和version,用于反序列化使用

const char* getPluginType() const;

const char* getPluginVersion() const;
初始化、配置、销毁函数

初始化函数,在这个插件准备开始run之前执行。一般申请权值显存空间并copy权值

int initialize();

terminate函数就是释放initialize开辟的一些显存空间

void terminate();

释放整个plugin占用的资源

void destroy();

​ 配置configurePlugin这个插件op,判断输入和输出类型数量是否正确。官方还提到通过这个配置信息可以告知TensorRT去选择合适的算法(algorithm)去调优这个模型。但自动调优目前还没有尝试过,一般自己写的plugin执行代码都是定死的,所谓的调优步骤可能更多地针对官方的op。

void MyCustomPluginDynamic::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* outputs, int nbOutputs)
{
    assert(nbOutputs == 1);
    assert(nbInputs == 2);
    assert(mType == inputs[0].desc.type);
};

​ TensorRT调用此方法以判断pos索引的输入/输出是否支持inOut[pos].format和inOut[pos].type指定的格式/数据类型。如果插件支持inOut[pos]处的格式/数据类型,则插件可以使其结果取决于inOut[0…pos-1]中的格式/数据类型,该格式/数据类型将设置为插件支持的值。这个函数不需要检查inOut[pos + 1…nbInputs + nbOutputs - 1],pos的决定必须仅基于inOut[0…pos]。

bool MyCustomPlugin::supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) 
{
    // 假设有一个输入和一个输出
    assert(0 <= pos && pos < 2);
    const auto *in = inOut;
    const auto *out = inOut + nbInputs;
    switch(pos){
        case 0:
            return in[0].type == DataType::kFLOAT && in[0].format == nvinfer1::TensorFormat::kLINEAR;
        case 1:
            return out[0].type == in[0].type && out[0].format == nvinfer1::TensorFormat::kLINEAR;
    }
};
运行相关函数

1、获得plugin所需要的显存大小。最好不要在plugin enqueue中使用cudaMalloc申请显存

size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const{
    // 计算这个op前向过程中需要的中间显存数量
    size_t need_num;
    return need_num * sizeof(float);
};

2、inference函数

int enqueue(int batchSize, const void* const* inputs, void** outputs, void *workspace, cudaStream_t stream){
    // 假设这个fun是需要的中间变量,可以直接使用TensorRT开辟的显存空间
    fun = static_cast<float*>(workspace);
};

​ 需要注意的是,如果操作中需要一些分布在显存中的中间变量,可以通过传过来的指针参数workspace获取。默认写的.cu是fp32的,TensorRT在fp16运行模式下,运行到不支持fp16的插件op时,会自动切换到fp32模式,等插件op运行完再切换回来。

​ 可以设置max workspace,避免显存移除,并且可以显存复用。如果在前一层使用cudaMalloc申请了显存,下一层无法使用(具体原因需要明确)另外,权值一般不复用,所以权值不会放在workspace里面,并且会使用cudaMalloc进行申请。

静态shape里面N维度是可变的,小于max batchsize。
下面以官方给出的其中1个sample,lReluPlugin.cpp中的enqueue函数为例:

leakyRelu的公式如下:
在这里插入图片描述

int LReLU::enqueue(int batchSize, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
    
    
    const void* inputData = inputs[0];
    void* outputData = outputs[0];
    pluginStatus_t status = lReLUInference(stream, mBatchDim * batchSize, mNegSlope, inputData, outputData);
    return status; 
}

其对应的CUDA内核函数在lReLU.cu里,我在里面加了注释,方便大家理解:

template <unsigned nthdsPerCTA>

__launch_bounds__(nthdsPerCTA) __global__ void pReLUKernel(const int n, const float negativeSlope, const float* input, float* output)
{
    
    
    // blockIdx.x表示当前线程块在线程格里x维度上的索引;nthdsPerCTA即blockDim.x,表示当前线程块中x维度上所有线程的个数;
    // threadIdx.x表示当前线程在线程块里x维度上的索引;gridDim.x表示当前线程格中x维度上所有线程块的个数;
    // i += gridDim.x * nthdsPerCTA,代表步长为gridDim.x * nthdsPerCTA,即1个线程格里的所有线程数。
    for(int i = blockIdx.x * nthdsPerCTA + threadIdx.x; i < n; i += gridDim.x * nthdsPerCTA)
    {
    
    
        //negativeSlope就是系数阿尔法
        output[i] = input[i] > 0 ? input[i] : input[i] * negativeSlope;
    }
}

pluginStatus_t lReLUGPU(cudaStream_t stream, const int n, const float negativeSlope, const void* input, void* output)
{
    
    
    // 这个n就是控制leakyRelu输出个数的变量
    const int BS = 512;
    const int GS = (n + BS - 1) / BS;
    // <BS>是模板参数,表示使用的线程块大小,可以传给内核函数pReLUKernel()
    pReLUKernel<BS><<<GS, BS, 0, stream>>>(n, negativeSlope, (const float*) input, (float*) output);
    return STATUS_SUCCESS;
}

pluginStatus_t lReLUInference(cudaStream_t stream, const int n, const float negativeSlope, const void* input, void* output)
{
    
    
    return lReLUGPU(stream, n, negativeSlope, (const float*) input, (float *) output);
}
IPluginCreator相关函数

总览:

class MyCustomPluginCreator : public BaseCreator

{
public:
	MyCustomPluginCreator();
	~MyCustomPluginCreator() override = default;
	const char* getPluginName() const override;
	const char* getPluginVersion() const override;
	const PluginFieldCollection* getFieldNames() override;
	IPluginV2DynamicExt* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) override;
	IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;

private:
	static PluginFieldCollection mFC;
	static std::vector<PluginField> mPluginAttributes;
	std::string mNamespace;
}

获得pluginname和version,用于辨识creator

const char* getPluginName() const;

const char* getPluginVersion() const;

通过PluginFieldCollection去创建plugin

将op需要的权重和参数一个一个取出来,然后调用上文提到的第一个构造函数:

const nvinfer1::PluginFieldCollection* getFieldNames();

IPluginV2DynamicExt* MyCustomPlugin::createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc)
{
    int in_channel;
    std::vector<float> weight;
    std::vector<float> bias;
    const PluginField* fields = fc ->fields;
    for (int i = 0; i < fc ->nbFields; ++i)
    {
        const char* attrName = fields[i].name;
        if (!strcmp(attrName, "in_channel"))
        {
            ASSERT(fields[i].type == PluginFieldType::kINT32);
            in_channel = *(static_cast<const int32_t*>(fields[i].data));
        }
        else if (!strcmp(attrName, "weight"))
        {
            ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
            int size = fields[i].length;
            h_weight.reserve(size);
            const auto* w = static_cast<const float*>(fields[i].data);
            for (int j = 0; j < size; j++)
            {
                h_weight.push_back(*w);
                w++;
            }
        }
        else if(!strcmp(attrName, "bias"))
        {
            ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
            int size = fields[i].length;
            h_bias.reserve(size);
            const auto* w = static_cast<const float*>(fields[i].data);
            for (int j = 0; j < size; j++)
            {
                h_bias.push_back(*w);
                w++;
            }
        }
    }
    
    Weights weightWeights{DataType::kFLOAT, weights.data(), (int64_t) weight.size()};
    Weights biasWeights{DataType::kFLOAT, bias.data(), (int64_t) _bias.size()};
    
    MyCustomPlugin* obj = new MyCustomPlugin(in_channel, weightWeights, biasWeights);
    obj -> setPluginNamespace(mNamespace.c_str());
    return obj;
}

​ PluginFieldCollection是成员变量,也会作为getFieldNames成员函数的返回类型。PluginFieldCollection的主要作用是传递这个插件op所需要的权重和参数,在实际的engine推理过程中并不使用,而在parse中会用到(例如caffe2trt、onnx2trt)。

​ 当使用这些parse去解析这个op的时候,这个op的权重和参数会经历Models–>TensorRT engine–>TensorRT runtime这个过程。

​ 举个例子,在onnx-tensorrt中,用DEFINE_BUILTIN_OP_IMPORTER去注册op,然后通过parse解析onnx模型,根据注册好的op去一个个解析构建模型,假设定义的op为my_custom_op,在DEFINE_BUILTIN_OP_IMPORTER(my_custom_op)会这样实现:

DEFINE_BUILTIN_OP_IMPORTER(my_custom_op)
{
	ASSERT(inputs.at(0).is_tensor(), ErrorCode::kUNSUPPORTED_NODE);
    ...
    const std::string pluginName = "CUSTOM-OP";
    const std::string pluginVersion = "001";
    
    // f保存这个op需要的权重和参数,从onnx模型中获取
    std::vector<nvinfer1::PluginField>f;
    f.emplace_back("in_channel", &in_channel, nvinfer1::PluginFieldType::kINT32, 1);
    f.emplace_back("weight", kernel_weights.values, nvinfer1::PluginFieldType::kFLOAT32, kernel_weights.count());
    f.emplace_back("bias", bias_weights.values, nvinfer1::PluginFieldType::kFLOAT32, bias_weights.count);
    
    // 从plugin工厂中获取该插件,并且将权重和参数传递进去
    nvinfer1::IPluginV2* plugin = importPluginFromRegistry(ctx, pluginName, pluginVersion, node.name(), f);
    RETURN_FIRST_OUTPUT(ctx->network()->addPluginV2(tensors.data(), tensors.size(), *plugin));
}

​ 进入importPluginFromRegistry函数内部,可以发现参数经过fc变量通过createPlugin传递给了plugin:

nvinfer1::IPluginV2* importPluginFromRegistry(IImporterContext* ctx, const std::string& pluginName,  const std::string& pluginVersion, const std::string& nodeName, const std::vector<nvinfer1::PluginField>& pluginFields)
{
    const auto mPluginRegistry = getPluginRegistry();
    const auto pluginCreator = mPluginRegistry->getPluginCreator(pluginName.c_str(), pluginVersion.c_str(),"ONNXTRT_NAMESPACE");
    
    if(!pluginCreator)
    {
        return nullptr;
    }
    // 接受传进来的权重和参数信息,传递给plugin
    nvinfer1::PluginFieldCollection fc;
    fc.nbFields = pluginFields.size();
    fc.fields = pluginFields.data();
    return pluginCreator->createPlugin(nodeName.c_str(), &fc);
}

​ 上述步骤中,会提供pluginName和pluginVersion初始化MyCustomPluginCreator,其中createPlugin成员函数是我们需要编写的。

创建一个空的mPluginAttributes初始化mFC:

MyCustomPluginCreator::MyCustomPluginCreator()
{
	mPluginAttributes.emplace_back(PluginField("in_channel", nullptr, PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(PluginField("weights", nullptr, PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kFLOAT32, 1));
    
    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

反序列化,调用反序列化那个构造函数,生成plugin

nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength);

​ 这个函数会被onnx-tensorrt的一个叫做TRT_IPluginV2的转换op调用,这个op会读取onnx模型的data数据将其反序列化到network中。

​ 这里的name和前面的type是一样的;这里的createPlugin需要自己写创建内容,然后自己去调用。Create函数可以用来封装接口并向外界提供库,不对外暴露plugin的设计,本质是通过传递结构体参数,利用其调用构造函数。

Dynamic Shape Plugin API

跟static shape相比有差异的函数

static implicit(隐式)batch vs dynamic explicit(显式) batch

1、根据输入个数和动态输入维度,获得第index个输出的动态维度

static

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

dynamic

nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder);

2、enqueue和getWorkspaceSize多了输入输出的信息、维度类型等

static

int enqueue(int batchSize, const void* const* inputs, void** outputs, void *workspace, cudaStream_t stream);

dynamic

int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream);

隐式和显式只会在plugin里面遇到。

静态shape的隐式batch意思是,这个batch的数值是enqueue传递进来的,剩下的维度都是确定,batch是动态的。静态shape中,TRT的推理中,batch是拿不到的,getOutputDimensions的inputs参数只会有CHW,是一个明确数值和维度的数组。对于enqueue函数,都是明确数值。

动态shape的显式batch是在getOutputDimensions函数中,inputs参数里面是NCHW,这几个维度的值都有。动态shape的输入维度数值都是不确定的,而输入输出之间的关系是通过exprBuilder来确定的,相当于一个四则运算器,做shape infer。对于enqueue函数,由于都是不确定的数值,需要输入输出的描述

静态是shape信息是可以提前拿到的,而动态只有在运行的时候才能获得的。

PluginCreator注册

​ 在加载NvInferRuntimeCommon.h头文件时,会得到一个getPluginRegistry,这里类中包含了所有已经注册了的IPluginCreator,在使用的时候通过getPluginCreator函数得到相应的IPluginCreator。

有两种注册方式:

1、调用API进行注册

extern "C" {
    bool initLibNvInferPlugins(void* logger, const char* libNamespace)
    {
        initializePlugin<nvinfer1::plugin::GridAnchorPluginCreator>(logger, libNamespace);
        initializePlugin<nvinfer1::plugin::NMSPluginCreator>(logger, libNamespace);
        initializePlugin<nvinfer1::plugin::ReorgPluginCreator>(logger, libNamespace);
        ...
        return true;
    }
} 

​ 其中initializePlugin函数执行了addPluginCreator函数:

template <typename CreatorType>
void initializePlugin(void* logger, const char* libNamespace)
{
    PluginCreatorRegistry::getInstance().addPluginCreator<CreatorType>(logger, libNamespace);
}

​ addPluginCreator函数又执行了getPluginRegistry() -> registerCreator对pluginCreator进行了注册,这样就完成注册任务了:

void addPluginCreator(void* logger, const char* libNamespace)
{
	...
        if(mRegistryList.find(pluginType) == mRegistryList.end())
        {
            bool status = getPluginRegistry()->registerCreator(*pluginCreator, libNamespace);
            if (status)
            {
                mRegistry.push(std::move(pluginCreator));
                mRegistryList.insert(pluginType);
                verboseMsg = "Plugin creator registration succeeded - " + pluginType;
            }
            else
            {
                errorMsg = "Could not register plugin creator: " + pluginType;
            }
        }
    	else
        {
            verboseMsg = "Plugin creator already registered - " + pluginType;
        }
   ...
}

2、直接通过REGISTER_TENSORRT_PLUGIN来注册:

// 在加载'NvinferRuntimeCommon.h'头文件的时候会得到一个'getPluginRegistry'
extern "C" TENSORRTAPI nvinfer1::IPluginRegistry* getPluginRegistry();
namespace nvinfer1
{
template <typename T>
class PluginRegistrar{

public:

	PluginRegistrar() {getPluginRegistry()->registerCreator(instance, "");}

private:

	T instance{};

};

#define REGISTER_TENSORRT_PLUGIN(name) \
static nvinfer1::PluginRegistrar<name> pluginRegistrar##name {}
}

​ 也就是说,如果我们已经在plugin的.h文件中执行了REGISTER_TENSORRT_PLUGIN(BatchedNMSPluginCreator); 就不需要再创建一个类似于官方的initLibNvInferPlugins()函数去一个一个注册了。

加载TRT库的时候,会得到一个全局变量,程序启动的时候,就会注册到全局变量。

如何使用注册好的PluginCreator?

class IPluginRegistry{

public:

	virtual bool registerCreator(IPluginCreator& creator, const char* pluginNamespace) noexcept = 0;

	virtual IPluginCreator* const* getPluginCreatorList(int* numCreators) const noexcept = 0;

	virtual IPluginCreator* getPluginCreator(const char* pluginType, const char* pluginVersion, const char* pluginNamespace = "") noexcept = 0;

}

TRT在线上推理的时候,从文件里读反序列化的模型时,可以读到pluginType和pluginVersion,再去调用getPluginCreator,通过这两个字符串,获得自己写的pluginCreator,再调用这个Creator的反序列化函数,就可以从文件中把权值和参数反序列化为一个plugin,再调用plugin的enqueue函数去执行推理。这就是注册的原因,只有注册了,线上才能自动反序列化去执行plugin,否则就找不到。

如何调用自己写好的plugin?

​ 使用addPluginV2函数,例如:

IPluginV2Layer* embLayer = network->addPluginV2(inputs, 3, embPlugin);

TensorRT如何debug-debug plugin?

TRT是闭源软件,API相对比较复杂。

什么情况下要debug TensorRT Plugin?

1)无论是使用API还是parser构建网络,模型转换完后,结果误差很大;

2)增加了自定义plugin实现算子合并,结果对不上;

3)使用FP16或INT8优化策略后,算法精确度掉了很多。

规整网络和训练的好的网络适合做int8优化,其余的可能精度掉很多。

推荐几种debug方法:

1、使用parser转换网络,使用dump API接口,查看网络结构是否对的上;

2、使用了plugin,要写单元测试;

3、通用方法,打印输出:

1)官方建议:将可疑层的输出设置为network output(比较累);

2)作者的方法:增加一个debug plugin。链接:https://github.com/LitLeo/TensorRT_Tutorial/tree/master/resource_for_billibilli/debug_plugin

TensorRT Plugin转换成FP16:

​ 如果网络中没有Plugin,则直接调用下列代码即可实现FP32->FP16的转换:

config->setFlag(BuilderFlag::kFP16);

builder->platformHasFastFp16();          // 判断平台是否支持FP16

builder->platformHasFastInt8();           // 判断平台是否支持Int8 

​ 如果网络中有Plugin,则需要注意以下事项:

1)编写Plugin时需要注意的是:

(1)Enqueue函数要增加half版本;

(2)注意supportsFormatCombination函数。保证输入输出类型一致,并要求输入输出类型与mType一致。

2)fp16模型,输入设置为float类型还是half类型?

​ 都行,但建议是将输入设置成float。

3)模型要配合混合精度训练,否则可能会出现溢出问题。

代码示例:https://github.com/NVIDIA/TensorRT/blob/7.2.1/plugin/skipLayerNormPlugin/skipLayerNormPlugin.cpp

参考链接:

https://zhuanlan.zhihu.com/p/567244140

https://www.bilibili.com/video/BV19Y411g7YY/?spm_id_from=333.999.0.0&vd_source=8002c1ea19b925cd4fa92e8ddf798043

https://zhuanlan.zhihu.com/p/297002406

https://github.com/nvidia/TensorRT

猜你喜欢

转载自blog.csdn.net/sinat_41886501/article/details/129624918