说明
- 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);
}