一文玩转pytorch转onnx-tensorRT ——（A）onnx转tensorRT

说明

• onnx和tensorRT是分开的，onnx像是prototxt和weight的打包在一起的东西。所以由onnx转到tensorRT下，还需要让onnx能搜索到（或parsing）所对应的层。
• caffeparsing有注册自定义层的函数，而onnx没有，下面会就onnx-tensorRT的源码详细说到。对于自定义的onnx层，需要修改后重新编译onnx-tensorRT，查到的资料也是这样说
• 另外，在正式开始前，贴几个帖子，但是帮助不是特别大：Onnx-tensorrt详解之nvonnxparser库；以yolov3为例子，手写定义onnx的格式；找到的demo–example for custom layer

修改及编译onnx-tensorRT

• 版本：onnx和tensorRT肯定是对应的，所以要注意版本的问题。我现在使用的tensorRT说5.1，选用的onnx-tensorRT也是5.1，链接如上。下面说的东西都是针对这个工程了。

• 主调函数，看此工程cmakeFIle的话，可以看到多个动态链接库打包在一起的有PLUGIN_SOURCES IMPORTER_SOURCES RUNTIME_SOURCES ONNXIFI_SOURCES。但是一般调用是从nvonnxparser::createParser开始的，此函数位于onnxparser里，其实主要调用打包了modelImporter类。下面就针对ModelImporter.cpp（针对v5.1）来做简要说明：

• 函数的数据流：[函数所在行]，数据流所用函数名所在行

  • 1）parseFromFile[226]273;
  • 2)parse[449]453;
  • 3)parseWithWeightDescriptors[406]423; （这是直接返回的函数，可暂时不管）
  • 4)importModel[457]497;
  • 5)importNode[138].
  • 需要留意的是，在数据流到5）开始检查已经注册了的operations，而这在字典_op_importers里面。据此查找ModelImporter.hpp时，是能发现一个注册函数的尸体registerOpImporter，上面这么写道：Note: This allows existing importers to be replaced如果这个函数活着的话，它一定是在NvOnnxParser.h的一个虚函数（实际上在此头文件中没有），因为ModelImporter继承于它，而且还是对外开放开发的窗口。

• 修改添加自定义层
  这儿应该是由两种方式的：

  • 1）复活上面的注册函数，需要在NvOnnxParser.h添加虚函数作为接口，然后找一个地方实现改函数，另外在registerOpImporter的旁边还躺着两个尸体set，如若实现，需要看看是不是需要提供额外的信息。用这种方式的好处是——只需编译一次就可以了，以后就直接添加注册。
  • 2）如果因为麻烦不想使之复活，就需要在builtin_op_importers.cpp中使用DEFINE_BUILTIN_OP_IMPORTER添加东西了。先说一下DEFINE_BUILTIN_OP_IMPORTER，它是一个宏，刚开始差点看懵，其实它正常的操作是这样的。
    • a）函数a的声明；
    • b）调用写入字典的函数写入字符串和刚声明的函数；
    • c）函数的定义（但是不包括大括号）。
    • 这个宏就是把abc通用的代码都定义成了宏定义，下面是一堆堆函数体，第一次看到这个文件时，别像我一样一脸懵逼就好。剩下的就是仿写了，例子满屏幕都是。像我添加的GN层，由pytorch转过来，名字是ATen（这儿的ATen还可以修改成普通的函数名，不过放到以后吧）

      graph(%input.1 : Float(1, 4, 2, 8),
      	      %gn.weight : Float(4),
      	      %gn.bias : Float(4)):
      	%3 : Float(1, 4, 2, 8) = onnx::LeakyRelu[alpha=0.01](%input.1), scope: gnSOLE/LeakyReLU[lr]
      	%4 : Float(1, 4, 2, 8) = onnx::ATen[cudnn_enabled=1, eps=1e-05, num_groups=2, operator="group_norm"](%3, %gn.weight, %gn.bias), scope: gnSOLE/GroupNorm[gn]
      	return (%4)	
      

• 编译onnx-tensorRT

  • 编译到时候需要修改的内容：

    • 1）HEADERS中添加NvOnnxParserTypedefs，也可不添加，等运行程序报错的时候记得这儿有这么一回事，主要编译的时候不用再专门给此文件的路径了，直接和onnxparser同一个文件夹。（至于HEADERS是啥，往下看就明白了，在cmakefile里面）

      set(HEADERS
      NvOnnxParser.h
      NvOnnxParserRuntime.h
      NvOnnxParserTypedefs.h)
      

    • 2）在add_library(nvonnxparser_plugi所对应的include_directories中添加额外的头文件路径 /usr/local/cuda/include

    • 需要注意的问题：onnx-tensorRT有很多第三方软件包含：pytorch、tensorRT、mx等。所以安装等时候要用locate onnxParser搜索一遍路径，处理一下它们；或者在编译onnx之后，注意调用onnx库等时候，给对正确的路径，因为忽略了路径上的一个子文件夹，onnx parsing时还是找不到自定义的层，我在这里耽误了小半天。

• 那最后就上代码吧，在builtin_op_importers.cpp中所添加的代码。注意，pytorch转onnx时，有很多的ATen族的函数，因为我用到的只有一个，所以就没管。如果因为这个出了问题，可以使用switch或者在pytorch转onnx时进行修改。

DEFINE_BUILTIN_OP_IMPORTER(ATen) {
    ASSERT(inputs.at(0).is_tensor(), ErrorCode::kUNSUPPORTED_NODE);
    ASSERT(inputs.at(1).is_weights(), ErrorCode::kUNSUPPORTED_NODE);
    ASSERT(inputs.at(2).is_weights(), ErrorCode::kUNSUPPORTED_NODE);
    nvinfer1::ITensor &tensor = inputs.at(0).tensor();
    auto scale_weights = inputs.at(1).weights();
    auto bias_weights = inputs.at(2).weights();
    OnnxAttrs attrs(node);
    int cudnn_enabled = attrs.get<int>("cudnn_enabled", 1);
    float eps = attrs.get<float>("eps", 1e-5f);
    int num_groups = attrs.get<int>("num_groups", 32);
    std::string oper = attrs.get<std::string>("operator", ""), version = "1";
    // TODO: Check if ONNX "spatial" attribute is important (maybe changes mean and variance broadcasting?)
    ASSERT(scale_weights.type == ::ONNX_NAMESPACE::TensorProto::FLOAT &&
           bias_weights.type == ::ONNX_NAMESPACE::TensorProto::FLOAT,
           ErrorCode::kUNSUPPORTED_NODE);
    nvinfer1::Dims dims = tensor.getDimensions();
    int nchan = dims.d[0];
    nvinfer1::Dims weights_shape{1, {nchan}};
    ASSERT(scale_weights.shape == weights_shape, ErrorCode::kINVALID_NODE);
    ASSERT(bias_weights.shape == weights_shape, ErrorCode::kINVALID_NODE);
    int nweight = nchan;

    std::vector <nvinfer1::PluginField> Attr;
    Attr.emplace_back(nvinfer1::PluginField("count", &nweight, nvinfer1::PluginFieldType::kINT32, 1));
    Attr.emplace_back(nvinfer1::PluginField("num_groups", &num_groups, nvinfer1::PluginFieldType::kINT32, 1));
    Attr.emplace_back(nvinfer1::PluginField("eps", &eps, nvinfer1::PluginFieldType::kFLOAT32, 1));
    Attr.emplace_back(nvinfer1::PluginField("w", scale_weights.values,
                                            nvinfer1::PluginFieldType::kFLOAT32, nweight));
    Attr.emplace_back(nvinfer1::PluginField("b", bias_weights.values,
                                            nvinfer1::PluginFieldType::kFLOAT32, nweight));
    nvinfer1::PluginFieldCollection mFC = {int(Attr.size()), Attr.data()};

    auto creator = getPluginRegistry()->getPluginCreator(oper.c_str(), version.c_str());
    nvinfer1::IPluginV2 *gnPlugin = creator->createPlugin("gn", &mFC);
    nvinfer1::ITensor *ipt[] = {&tensor};

    // Fold the weights together into a single bias and scale
    for (int i = 0; i < nweight; ++i) {
        float scale = (static_cast<float const *>(scale_weights.values))[i];
        float bias = (static_cast<float const *>(bias_weights.values))[i];
        std::cout << i << "w: " << scale << "b: " << bias << std::endl;
    }
    auto layer = ctx->network()->addPluginV2(ipt, 1, *gnPlugin);
    ASSERT(layer, ErrorCode::kUNSUPPORTED_NODE);
    RETURN_FIRST_OUTPUT(layer);
}