【tensorrt】——双线性上采样插件（提供源码）

本文主要是介绍【tensorrt】——双线性上采样插件（提供源码），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

在这里插入图片描述

简介：
如果用nvidia的gpu，在推理的时候，采用tensorrt进行加速是一个很好的选择，虽然tensorrt没有开源。

我一般选择的模型训练到部署的流程是：

pytorch训练模型
onnx模型导出
onnx模型转ncnn，mnn，tensorrt等模型
嵌入式推理框架，推理脚本书写。

这里用tensorrt做语义分割网络pspnet的推理加速。技术路线采用：pytorch——onnx——tensorrt。

1. pytorch——onnx

pytorch是内嵌了onnx模型导出的。这里pytorch版本的选择由使用的tensorrt的版本确定。这里我们采用TensorRT-YOLOv4项目中onnx-tensorrt中的tensorrt版本5.1xx。

这个版本上采样onnx中还是upsample，对应到pytorch<=1.0。pytorch1.0是支持nearest，bilinear两种方式的导出的。

2. onnx——tensorrt

TensorRT-YOLOv4中有resizenearest插件是没有双线性插值的。

2.1 写插件

resizebilinear是没有网络权重参数的，所以没有序列化重构，可以需要对以下进行重构。
需要重构：

getPluginType：
getOutputDimensions：计算网络输出tensor的尺寸
initialize：
enqueue：前向推理的具体入口

#pragma once#include "plugin.hpp"
#include "serialize.hpp"
#include <cassert>class ResizeBilinearPlugin final : public onnx2trt::Plugin {int   _ndims;float _scale[nvinfer1::Dims::MAX_DIMS];nvinfer1::Dims _output_dims;protected:void deserialize(void const* serialData, size_t serialLength) {deserializeBase(serialData, serialLength);deserialize_value(&serialData, &serialLength, &_ndims);deserialize_value(&serialData, &serialLength, &_scale);}size_t getSerializationSize() override {return serialized_size(_ndims) + serialized_size(_scale) + getBaseSerializationSize();}void serialize(void *buffer) override {serializeBase(buffer);serialize_value(&buffer, _ndims);serialize_value(&buffer, _scale);}public:ResizeBilinearPlugin(std::vector<float> const& scale): _ndims(scale.size()) {assert(scale.size() <= nvinfer1::Dims::MAX_DIMS);std::copy(scale.begin(), scale.end(), _scale);}ResizeBilinearPlugin(void const* serialData, size_t serialLength) {this->deserialize(serialData, serialLength);}virtual const char* getPluginType() const override { return "ResizeBilinear"; }virtual int getNbOutputs() const override { return 1; }virtual nvinfer1::Dims getOutputDimensions(int index,const nvinfer1::Dims *inputs, int nbInputDims) override;virtual int initialize() override;int enqueue(int batchSize,const void *const *inputs, void **outputs,void *workspace, cudaStream_t stream) override;
};

重构之后，一般都会向，tensorrt怎么调了。具体可以参考：【onnx-tensorrt】——源码阅读记录

总结就是：你看不到调用的接口，你只能模仿着写。

2.2 注册插件

在 builtin_plugins.cpp 中注册插件

REGISTER_BUILTIN_PLUGIN("FancyActivation",       FancyActivationPlugin);        // 相当于入库
REGISTER_BUILTIN_PLUGIN("ResizeNearest",         ResizeNearestPlugin);
REGISTER_BUILTIN_PLUGIN("ResizeBilinear",        ResizeBilinearPlugin);
REGISTER_BUILTIN_PLUGIN("Split"        ,         SplitPlugin);
REGISTER_BUILTIN_PLUGIN("InstanceNormalization", InstanceNormalizationPlugin);
REGISTER_BUILTIN_NVPLUGIN("Concat", ConcatPlugin);
REGISTER_BUILTIN_PLUGIN("DCNv2", DCNv2Plugin);
REGISTER_BUILTIN_PLUGIN("Mish", MishPlugin);
REGISTER_BUILTIN_PLUGIN("YOLO", YOLOPlugin);
REGISTER_BUILTIN_PLUGIN("DarkNetAdd", ADDPlugin);

注意：
注册插件的字符串ResizeBilinear和 virtual const char* getPluginType() const override { return “ResizeBilinear”; }的字符串保持一致。

2.3 使用插件，修改builtin_op_importers.cpp

插件写好了，什么时候使用的呢？我怎么让tensorrt使用我的插件呢？

答案： 具体是在builtin_op_importers.cpp中进行控制的，这里以upsample为例子：

DEFINE_BUILTIN_OP_IMPORTER(Upsample) {ASSERT(inputs.at(0).is_tensor(), ErrorCode::kUNSUPPORTED_NODE);nvinfer1::ITensor &tensor = inputs.at(0).tensor();ASSERT(tensor.getDimensions().nbDims == 3, ErrorCode::kUNSUPPORTED_NODE);OnnxAttrs attrs(node);float height_scale, width_scale;if (ctx->getOpsetVersion() >= 9) {ASSERT(inputs.size() == 2, ErrorCode::kINVALID_NODE);auto scales_input = inputs.at(1);ASSERT(scales_input.is_weights(), ErrorCode::kUNSUPPORTED_NODE);ShapedWeights scales_weights = scales_input.weights();ASSERT(scales_weights.shape.nbDims == 1, ErrorCode::kUNSUPPORTED_NODE);ASSERT(scales_weights.count() == 4, ErrorCode::kUNSUPPORTED_NODE);ASSERT(scales_weights.type == ::ONNX_NAMESPACE::TensorProto::FLOAT,ErrorCode::kINVALID_NODE);float const *scales_ptr = static_cast<float const *>(scales_weights.values);ASSERT(scales_ptr[0] == 1 && scales_ptr[1] == 1,ErrorCode::kUNSUPPORTED_NODE);height_scale = scales_ptr[2];width_scale = scales_ptr[3];} else {if (!attrs.count("scales")) {height_scale = attrs.get<float>("height_scale");width_scale = attrs.get<float>("width_scale");} else {auto scales = attrs.get<std::vector<float>>("scales");ASSERT(scales.size() == 4, ErrorCode::kUNSUPPORTED_NODE);ASSERT(scales[0] == 1 && scales[1] == 1, ErrorCode::kUNSUPPORTED_NODE);height_scale = scales[2];width_scale = scales[3];}}auto scale = {height_scale, width_scale};auto mode = attrs.get<std::string>("mode", "nearest");        // 默认采用 nearest 上采样方式ASSERT(mode == "nearest" || "linear", ErrorCode::kUNSUPPORTED_NODE);if (mode == "nearest")RETURN_FIRST_OUTPUT(ctx->addPlugin(new ResizeNearestPlugin(scale), {&inputs.at(0).tensor()}));        // 这里确定使用何种自定义的类别插件else if (mode == "linear")RETURN_FIRST_OUTPUT(ctx->addPlugin(new ResizeBilinearPlugin(scale), {&inputs.at(0).tensor()}));}

DEFINE_BUILTIN_OP_IMPORTER(Upsample)是在onnx模型解析导入的时候调用的。
ctx->addPlugin(new ResizeBilinearPlugin(scale), {&inputs.at(0).tensor()})就是初始化一个类，后续tensorrt模型的序列化，推理就会使用新定义的（自己定义的）网络层

2 动态库编译

前面插件也写好了，onnx模型也能解析了，别人怎么用呢？

参考TensorRT-YOLOv4直接给源码编译
编译成动态库，给别人动态库，别人直接用动态库

这里还是使用TensorRT-YOLOv4编译成动态库。
其实动态库中已经包含了我们刚才修改的文件：

resizebilinear.cu
resizebilinear.h
builtin_op_importers.cpp

后续在使用的时候还是用tensorrt原有的头文件，链接时候，链接上前面编译好的动态库就好了。比如：

cmake_minimum_required(VERSION 2.8)
project(net)find_package(CUDA REQUIRED)include_directories(../include   
)set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast ")
set(CUDA_NVCC_FLAGS  "-D_FORCE_INLINES -Xcompiler -fPIC -gencode arch=compute_${GPU_ARCHS},code=sm_${GPU_ARCHS} -gencode arch=compute_${GPU_ARCHS},code=compute_${GPU_ARCHS}")# packed so library
set(srcs net.cpp resize.cu)
cuda_add_library(megengine SHARED ${srcs})
target_link_libraries(megengine mynvonnxparser              # 不用包含tensorrt plugin层的头文件，采用原有的头文件就可以mynvonnxparser_runtime)# 1. 不需要包含opencv_libs， 因为没有使用opencv的操作