Tensort 8.5.1.7 注册自定义算子(去掉mmdeploy框架)

最新推荐文章于 2024-01-06 21:55:48 发布

knighthyyyyyy

最新推荐文章于 2024-01-06 21:55:48 发布

阅读量535

点赞数 2

文章标签：人工智能自动驾驶

本文链接：https://blog.csdn.net/qq_34928047/article/details/130648090

版权

注册自定义算子流程

1.准备工作

下载对应的Tensor-SDK
在这里插入图片描述
TensorRT-8.5.1.7内部详情(解压之后文件夹详情如下)

2.克隆官方开源代码，进行编译

https://github.com/NVIDIA/TensorRT.git
# cd 进入TensorRT

#切换分支
git checkout 1d6bf36

# 克隆全部分支 一般 onnx 、caffe 、common
git submodule update --init --recursive

3.更新bashrc

export PATH=/home/ying/data/code/TensorRT/TensorRT-8.5.1.7/include${PATH:+:${PATH}}
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/ying/data/code/TensorRT/TensorRT-8.5.1.7/lib


#更新bashrc source ~/.bashrc

4.注册

将写好的文件放入plugin 路径下
新增bevpoolv2
编辑plugin文件下的CMakeLists 在68行新增注册算子名称

set(PLUGIN_LISTS
    batchTilePlugin
    batchedNMSPlugin
    clipPlugin
    coordConvACPlugin
    cropAndResizePlugin
    decodeBbox3DPlugin
    detectionLayerPlugin
    disentangledAttentionPlugin
    efficientNMSPlugin
    efficientNMSPlugin/tftrt
    flattenConcat
    generateDetectionPlugin
    gridAnchorPlugin
    groupNormalizationPlugin
    instanceNormalizationPlugin
    leakyReluPlugin
    multilevelCropAndResizePlugin
    multilevelProposeROI
    multiscaleDeformableAttnPlugin
    nmsPlugin
    normalizePlugin
    nvFasterRCNN
    pillarScatterPlugin
    priorBoxPlugin
    proposalLayerPlugin
    proposalPlugin
    pyramidROIAlignPlugin
    regionPlugin
    reorgPlugin
    resizeNearestPlugin
    roiAlignPlugin
    scatterPlugin
    specialSlicePlugin
    splitPlugin
    voxelGeneratorPlugin
    bevpoolv2
    )

修改plugin/apiInferPlugin.cpp 在50行修改：
在这里插入图片描述
文件最后加入插入图片描述
修改parsers/onnx/builtin_op_importers.cpp 新增算子


DEFINE_BUILTIN_OP_IMPORTER(bev_pool_v2)
{

ASSERT(inputs.at(0).is_tensor(), ErrorCode::kUNSUPPORTED_NODE);
ASSERT(inputs.at(1).is_tensor(), ErrorCode::kUNSUPPORTED_NODE);
ASSERT(inputs.at(2).is_weights(), ErrorCode::kUNSUPPORTED_NODE);
ASSERT(inputs.at(3).is_weights(), ErrorCode::kUNSUPPORTED_NODE);
ASSERT(inputs.at(4).is_weights(), ErrorCode::kUNSUPPORTED_NODE);
ASSERT(inputs.at(5).is_weights(), ErrorCode::kUNSUPPORTED_NODE);
ASSERT(inputs.at(6).is_weights(), ErrorCode::kUNSUPPORTED_NODE);

std::vector<nvinfer1::ITensor*> tensors;
for (auto& input : inputs)
 {
        auto* tensorPtr = &convertToTensor(input, ctx);
        tensors.push_back(tensorPtr);
 }
//std::vector<nvinfer1::ITensor*> tensors {&inputs.at(0).tensor(), &inputs.at(1).tensor(), &inputs.at(2).tensor(), &inputs.at(3).tensor(), 
//&inputs.at(4).tensor(), &inputs.at(5).tensor(), &inputs.at(6).tensor()};
const std::string pluginName = "bev_pool_v2";
const std::string pluginVersion = "1";
OnnxAttrs attrs(node, ctx);

std::vector<nvinfer1::PluginField> f;
int output_height,output_width;
output_height=128;
output_width=128;
f.emplace_back("output_height", &output_height, nvinfer1::PluginFieldType::kINT32, 1);
f.emplace_back("output_width",  &output_width, nvinfer1::PluginFieldType::kINT32, 1);

auto const  plugin = createPlugin(getNodeName(node), importPluginCreator(pluginName, pluginVersion), f); 
ASSERT(plugin != nullptr && "bev_pool_v2 plugin was not found in the plugin registry!",
        ErrorCode::kUNSUPPORTED_NODE);


//std::cout<<"tensors.size() :: "<<tensors.size()<<std::endl;
RETURN_FIRST_OUTPUT(ctx->network()->addPluginV2(tensors.data() ,tensors.size(), *plugin));
//std::cout<<"ctx->network()->addPluginV2---->end "<<std::endl;

}