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

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

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

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

这里用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的操作

说明：

• 这里的mynvonnxparser， mynvonnxparser_runtime就是前面编译的动态库
• tensorrt的头文件我放到了/usr/include下，所有cmake中没有指定

other

• 下载地址
• 一定要主要tensorrt的版本，不同版本插件的书写是不一样的
  没有积分的留下邮箱吧