实现TensorRT-7.0插件自由!(如果不踩坑使用TensorRT插件功能) (3)

实际插件op的执行函数,我们自己实现的cuda操作就放到这里(当然C++写的op也可以放进来,不过因为是CPU执行,速度就比较慢了),与往常一样接受输入inputs产生输出outputs,传给相应的指针就可以。

int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream){ // 假如这个fun是你需要的中间变量 这里可以直接用TensorRT为你开辟的显存空间 fun = static_cast<float*>(workspace); }

需要注意的是,如果我们的操作需要一些分布在显存中的中间变量,可以通过传过来的指针参数workspace获取,上述代码简单说明了一下使用方法。

再多说一句,我们默认写的.cu是fp32的,TensorRT在fp16运行模式下,运行到不支持fp16的插件op时,会自动切换到fp32模式,等插件op运行完再切换回来。

getOutputDimensions

TensorRT支持Dynamic-shape的时候,batch这一维度必须是explicit的,也就是说,TensorRT处理的维度从以往的三维[3,-1,-1]变成了[1,3,-1,-1]。最新的onnx-tensorrt也必须设置explicit的batchsize,而且这个batch维度在getOutputDimensions中是可以获取到的。

在旧版的IPluginV2类中,getOutputDimensions的定义如下:

virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRTNOEXCEPT = 0;

而在新版的IPluginV2DynamicExt类中定义如下:

virtual DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) = 0;

我们要做的就是在这个成员函数中根据输入维度推理出模型的输出维度,需要注意的是,虽然说输出维度
是由输入维度决定,但这个输出维度其实“内定”的(也就是在计算之前就算出来了)。如果咱的插件op的输出维度需要通过实际运行计算得到,那么这个函数就无法满足咱了。

去你妈的好气哦

set/getPluginNamespace

为这个插件设置namespace名字,如果不设置则默认是"",需要注意的是同一个namespace下的plugin如果名字相同会冲突。

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(mycustom_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成员函数是我们需要编写的(下文会说)。

configurePlugin

配置这个插件op,判断输入和输出类型数量是否正确。官方还提到通过这个配置信息可以告知TensorRT去选择合适的算法(algorithm)去调优这个模型。

但自动调优目前还没有尝试过,我们一般自己写的plugin执行代码都是定死的,所谓的调优步骤可能更多地针对官方的op。

内容版权声明:除非注明,否则皆为本站原创文章。

转载注明出处:https://www.heiqu.com/wpwjpf.html