自定义plugin
流程
- 首先明确要开发的算子,最好是 CUDA 实现;
- 继承 IPluginV2DynamicExt / IPluginV2IOExt类实现一个Plugin 类,在这里调用前面实现的算子;
- 继承 IPluginCreator 类实现一个 PluginCreator 类,用于创建插件实例,然后注册该 Creator 类;
- 编译插件项目,生成动态链接库;
- 在构造 engine 之前,首先加载上一步编译出来的插件动态链接库,在构造 engine 时 TensorRT 会自动找到先前注册的插件创建器。
注意事项
Static Shape,用IPluginV2IOExt;Dynamic Shape,则使用IPluginV2DynamicExt。
IPluginV2
只能支持 implicit mode,
所以只能使用 execute
接口,并指定 batch:
builder->setMaxBatchSize(3);
.
.
.
context->execute(batch_size, buffers);
IPluginV2DynamicExt 支持动态shape,仅支持显示batch。
使用动态shape时,config中需要设置profile
IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions("input", OptProfileSelector::kMIN, Dims4(1, C, H, W));
profile->setDimensions("input", OptProfileSelector::kOPT, Dims4(2, C, H, W));
profile->setDimensions("input", OptProfileSelector::kMAX, Dims4(4, C, H, W));
config->addOptimizationProfile(profile);
最后运行时,需要设置实际运行shape:
context->setBindingDimensions(inputIndex, Dims4(3, 4, 2, 2));
这和 IPluginV2 demo中的隐式batch机制很类似,隐式batch需要优化时设置最大batch数,运行时需要设置实际的batch数目。
写自定义plugin时,推荐使用 IPluginV2DynamicExt做基类。支持静态/动态shape,显示batch也更直观。
IPluginV2的隐式batch模式下,plugin内部只能看到三维的shape信息,batch信息在enqueue函数内才能看到(context->execute(batch)传入。 IPluginV2DynamicExt 显示batch模式下,可以看到四维shape信息.
如果网络中有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
测试:
1)使用了plugin,要写单元测试;
2)使用parser转换网络,使用dump API接口,查看网络结构是否对的上
3)通用方法,打印输出:
a)官方建议:将可疑层的输出设置为network output(比较累);
b)自己写个debug plugin
函数解释
需要写2个类:
1)MyCustomPlugin,继承IPluginV2Ext/IPluginV2IOExt/IPluginV2DynamicExt,是插件类,用于写插件的具体实现;
2)MyCustomPluginCreator,继承BaseCreator, 负责创建和管理 MyCustomPlugin
实例,以及向TensorRT注册插件。
static Shape Plugin
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、用于查询本插件序列化需要的内存大小,实际上就是对所有当前类变量数据的字节大小求和。
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)去调优这个模型。
该方法用于对插件配置输入输出相关参数,且在 engine 构建阶段和执行阶段都会被调用,原因是构建阶段和执行阶段输入输出张量的维度信息可能不同(因为是 dynamic shape 的),因此需要在每次执行前都重新配置一下。
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
所指定张量的 type
和 format
的组合是否是被当前插件所支持的。type
无非就单精度、半精度、整型等等,而 format
则是指张量的布局方式
- pos 表示当前查询张量序号,注意这里输入和输出是合在一起排序的,也就是说
0 < pos < nbInputs + nbOutputs
,其中nbInputs
表示输入张量的个数,nbOutputs
表示输出张量的个数。当pos < nbInputs
时,表示当前查询的是输入张量,否则表示当前查询的是输出张量。- inOut 表示输入或输出张量的描述信息,其中包含了张量的维度信息,数据类型
type
,数据布局格式format
等。
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、插件执行方法,在这里调用 CUDA 算子。
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,避免显存移除,并且可以显存复用。
lReluPlugin.cpp中的enqueue函数为例:
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);
}
static shape IPluginCreator
class MyCustomPluginCreator : public BaseCreator
{
public:
MyCustomPluginCreator();
~MyCustomPluginCreator() override = default;
const char* getPluginName() const override;
const char* getPluginVersion() const override;
const PluginFieldCollection* getFieldNames() override;
// 通过PluginFieldCollection去创建plugin,将所需的参数和权值取出,调用MyCustomPlugin(args ...)
IPluginV2DynamicExt* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) override;
// 反序列化,调用MyCustomPlugin(const void* data, size_t length)来创建plugin
IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
private:
static PluginFieldCollection mFC;
static std::vector<PluginField> mPluginAttributes;
std::string mNamespace;
}
获得plugin name和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)
IPluginV2* createPlugin(const char* name,
const PluginFieldCollection* fc) noexcept override;
这是创建插件的主要方法,其中 name
表示插件名称,fc
表示插件类的字段集合,通过 fc -> fields
方法我们可以拿到 PluginField
指针数组,每个 PluginField
对象包含了字段名称,字段类型,字段数据等信息,通过类型转换可以得到具体的字段数据并创建插件实例。
IPluginV2* deserializePlugin(const char* name,
const void* serialData,
size_t serialLength) noexcept override;
该方法用于反序列化插件,其中 name
表示插件名称,serialData
表示序列化数据,serialLength
表示序列化数据的字节大小
Dynamic Shape Plugin API
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);
enqueue 在 TensorRT7 和 TensorRT8 中函数声明不同,升级 TRT 版本时注意调整
隐式和显式只会在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。
REGISTER_TENSORRT_PLUGIN注册
your_Plugin.cpp
REGISTER_TENSORRT_PLUGIN(GeluPluginDynamicCreator);
API注册
需要在plugin/api/InferPlugin.cpp里添加初始化plugin的接口:
1) 添加头文件
2)添加初始化插件的接口
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;
}
...
}
编译tensorrt
htop
ranger
tensorrt文档
源码中 doc/pdf
plugin例子和原理
https://zhuanlan.zhihu.com/p/297002406
demo
TensorRT-8.0.1.6/samples/sampleUffPluginV2Ext
在每个函数里增加log printf(),程序运行时可以每个函数运行顺序
cd TensorRT-8.0.1.6/samples/sampleUffPluginV2Ext
make
cd ../../bin/
./sample_uff_plugin_v2_ext | tee log.txt
找不到动态库libnvinfer.so.6,export 动态库
没有序列化和反序列化
ctreator
补充
build 时的log看出:
显性和隐性batch
TensorRT系列——explicit_batch vs implicit_batch_51CTO博客_tensorRt
显性batch
// 动态维度,设置batch = 4
context->setBindingDimensions(0, Dims4(4, 1, 112, 112));
context->executeV2(buffers);
隐性batch
For implicit batch, use createNetwork
or pass a 0 to createNetworkV2
.
builder = trt.Builder(...)
builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))