onnx学习

有点时间了，整理下部分学习资料，方便后续查找，先从onnx开始吧。(文章内容大部分摘取于网络资源）

1 onnx介绍

ONNX的本质，是一种Protobuf格式文件。onnx是基于protobuf来做数据存储和传输,*.proto后缀文件, 其定义是protobuf语法，类似json。

onnx-ml.proto描述onnx文件如何组成的，具有什么结构，里面定义了onnx的变量结构、类型等，其中几个核心部分是：ModelProto, GraphProto, NodeProto, ValueInfoProto, TensorProto,AttributeProto

onnx-ml.proto路径：https://github.com/onnx/onnx/blob/main/onnx/onnx-ml.proto

  • ModelProto:当加载了一个onnx后，会获得一个ModelProto。它包含一个GraphProto和一些版本，生产者的信息。
  • GraphProto: 包含了四个repeated数组(可以用来存放N个相同类型的内容，key值为数字序列类型.)。这四个数组分别是node(NodeProto类型)，input(ValueInfoProto类型)，output(ValueInfoProto类型)和initializer(TensorProto类型)；
  • NodeProto: 存node，放了模型中所有的计算节点,语法结构如下
  • ValueInfoProto: 存input，放了模型的输入节点。存output，放了模型中所有的输出节点；
  • TensorProto: 存initializer，放了模型的所有权重参数
  • AttributeProto:每个计算节点中还包含了一个AttributeProto数组，用来描述该节点的属性，比如Conv节点或者说卷积层的属性包含group，pad，strides等等；

用过protbuf的同学应该比较熟悉protobuf的使用流程（参考之前的文章：https://www.cnblogs.com/silence-cho/p/14544004.html）。

对于onnx，则是利用Protobuf的下面命令，将onnx-ml.proto编译得到c++使用的onnx-ml.pb.h和onnx-ml.pb.cc， 以及pyrhon使用的onnx_ml_pb2.py

protoc onnx-ml.proto --cpp_out=pbout --python_out=pbout
  • c++使用onnx-ml.pb.h和onnx-ml.pb.cc提供的函数，对onnx文件进行读取和修改
  • python程序使用onnx_ml_pb2.py提供的函数，对onnx文件进行读取和修改

最后总结下，onnx文件的结构如下图所示：

各模块的功能如下：

  • model：表示整个onnx的模型，包含图结构和解析器格式、opset版本、导出程序类型

  • model.graph：表示图结构，通常是我们netron看到的主要结构

  • model.graph.node：表示图中的所有节点，数组，例如conv、bn等节点就是在这里的，通过input、output表示节点之间的连接关系

  • model.graph.initializer：权重类的数据大都储存在这里

  • model.graph.input：整个模型的输入储存在这里，表明哪个节点是输入节点，shape是多少

  • model.graph.output：整个模型的输出储存在这里，表明哪个节点是输出节点，shape是多少

(对于anchorgrid类的常量数据，通常会储存在model.graph.node中，并指定类型为Constant，该类型节点在netron中可视化时不会显示出来)

2 onnx生成

python可以利用pytorch的torch.onnx.export()函数生成onnx文件，也可以利用onnx python包提供的接口生成onnx文件.

pytorch生成onnx

利用torch.onnx.export()函数，可以将模型转换成onnx文件， os.path.dirname(torch.onnx.__file__)包含了各个版本的onnx转换代码的逻辑，

示例代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import os

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(1, 1, 3, padding=1)
        self.relu = nn.ReLU()
        self.conv.weight.data.fill_(1)
        self.conv.bias.data.fill_(0)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

# 这个包对应opset11的导出代码，如果想修改导出的细节，可以在这里修改代码
# import torch.onnx.symbolic_opset11
print("对应opset文件夹代码在这里：", os.path.dirname(torch.onnx.__file__))

model = Model()
dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
    model, 

    # 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
    (dummy,), 

    # 储存的文件路径
    "demo.onnx", 

    # 打印详细信息
    verbose=True, 

    # 为输入和输出节点指定名称，方便后面查看或者操作
    input_names=["image"], 
    output_names=["output"], 

    # 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
    opset_version=11, 

    # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
    # 通常，我们只设置batch为动态，其他的避免动态
    dynamic_axes={
        "image": {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch", 2: "height", 3: "width"},
    }
)

print("Done.!")

onnx python包

通过import onnx和onnx.helper提供的make_node，make_graph，make_tensor等接口我们可以生成onnx文件，需要完成对node，initializer，input，output，graph，model的填充，代码如下：

import onnx # pip install onnx>=1.10.2
import onnx.helper as helper
import numpy as np

# https://github.com/onnx/onnx/blob/v1.2.1/onnx/onnx-ml.proto

nodes = [
    helper.make_node(
        name="Conv_0",   # 节点名字，不要和op_type搞混了
        op_type="Conv",  # 节点的算子类型, 比如'Conv'、'Relu'、'Add'这类，详细可以参考onnx给出的算子列表
        inputs=["image", "conv.weight", "conv.bias"],  # 各个输入的名字，结点的输入包含：输入和算子的权重。必有输入X和权重W，偏置B可以作为可选。
        outputs=["3"],  
        pads=[1, 1, 1, 1], # 其他字符串为节点的属性，attributes在官网被明确的给出了，标注了default的属性具备默认值。
        group=1,
        dilations=[1, 1],
        kernel_shape=[3, 3],
        strides=[1, 1]
    ),
    helper.make_node(
        name="ReLU_1",
        op_type="Relu",
        inputs=["3"],
        outputs=["output"]
    )
]

# 权重
initializer = [
    helper.make_tensor(
        name="conv.weight",
        data_type=helper.TensorProto.DataType.FLOAT,
        dims=[1, 1, 3, 3],
        vals=np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float32).tobytes(),
        raw=True
    ),
    helper.make_tensor(
        name="conv.bias",
        data_type=helper.TensorProto.DataType.FLOAT,
        dims=[1],
        vals=np.array([0.0], dtype=np.float32).tobytes(),
        raw=True
    )
]

inputs = [
    helper.make_value_info(
        name="image",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", 1, 3, 3]
        )
    )
]

outputs = [
    helper.make_value_info(
        name="output",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", 1, 3, 3]
        )
    )
]

graph = helper.make_graph(
    name="mymodel",
    inputs=inputs,
    outputs=outputs,
    nodes=nodes,
    initializer=initializer
)

# 如果名字不是ai.onnx，netron解析就不是太一样了
opset = [
    helper.make_operatorsetid("ai.onnx", 11)
]

# producer主要是保持和pytorch一致
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.9")
onnx.save_model(model, "my.onnx")

print(model)
print("Done.!")

3 读取和修改

由于protobuf任何支持的语言，我们可以使用[c/c++/python/java/c#等等]实现对onnx文件的读写操作。

onnx python包提供的接口，还可以用来对onnx文件进行读取，修改和保存。

  • 增：一般伴随增加node和tensor
graph.initializer.append(xxx_tensor)
graph.node.insert(0, xxx_node)

  • 删：

    graph.node.remove(xxx_node)

  • 改：

    input_node.name = 'data'

读取代码如下：

（注意：通过graph可以访问参数，数据是以protobuf的格式存储的，因此当中的数值会以bytes的类型保存。需要用np.frombuffer方法还原成类型为float32的ndarray。注意还原出来的ndarray是只读的。）

import onnx
import numpy as np

model = onnx.load("./my_model.onnx")
# print(helper.printable_graph(model.graph))
print(model)

print("==========input==========")
input = model.graph.input[0]
print("input: ", input.name)

print("==========node==========")
conv_node = model.graph.node[0]
relu_node = model.graph.node[1]
print(conv_node.name, conv_node.input, conv_node.output)
print(relu_node.name, relu_node.input, relu_node.output)

conv_weight = model.graph.initializer[0]
conv_bias = model.graph.initializer[1]

# 数据是以protobuf的格式存储的，因此当中的数值会以bytes的类型保存，通过np.frombuffer方法还原成类型为float32的ndarray
print(f"===================={conv_weight.name}==========================")
print(conv_weight.name, np.frombuffer(conv_weight.raw_data, dtype=np.float32))

print(f"===================={conv_bias.name}==========================")
print(conv_bias.name, np.frombuffer(conv_bias.raw_data, dtype=np.float32))


print("==========output==========")
output = model.graph.output[0]
print("output: ", output.name)

修改onnx文件中卷积权重值，代码如下：

import onnx
import onnx.helper as helper
import numpy as np

model = onnx.load("demo.onnx")

# 可以取出权重
conv_weight = model.graph.initializer[0]
conv_bias = model.graph.initializer[1]
# 修改权
conv_weight.raw_data = np.arange(9, dtype=np.float32).tobytes()

# 修改权重后储存
onnx.save_model(model, "demo.change.onnx")
print("Done.!")

4 合并两个onnx模型

可以将一个大模型拆成多个onnx小模型，也可以将多个onnx模型合并成一个大模型。下面是合并两个onnx模型的示例代码：

##################################### export onnx

import torch
from torchvision import models

model1 = models.resnet18(True).eval()
model2 = models.resnet50(True).eval()

input_tensor = torch.zeros(1, 3, 224, 224)
torch.onnx.export(
    model1, (input_tensor,), "model1.onnx", opset_version=11, input_names=["image"], output_names=["prob"]
)

torch.onnx.export(
    model2, (input_tensor,), "model2.onnx", opset_version=11, input_names=["image"], output_names=["prob"]
)

with torch.no_grad():
    print("Torch Out1 = ", model1(input_tensor)[0, :5].data.numpy())
    print("Torch Out2 = ", model2(input_tensor)[0, :5].data.numpy())


############################################## merge onnx

import onnx

model1 = onnx.load("model1.onnx")
model2 = onnx.load("model2.onnx")

def rename_model(model, newname):

    for n in model.graph.node:
        n.name = newname(n.name)
        for i in range(len(n.input)):
            n.input[i] = newname(n.input[i])

        for i in range(len(n.output)):
            n.output[i] = newname(n.output[i])

    for n in model.graph.initializer:
        n.name = newname(n.name)

    for n in model.graph.input:
        n.name = newname(n.name)

    for n in model.graph.output:
        n.name = newname(n.name)


def newname_func(prefix):
    def impl(name):
        return f"{prefix}-{name}"
    return impl

rename_model(model1, newname_func("model1"))
rename_model(model2, newname_func("model2"))

model1.graph.node.extend(model2.graph.node)
model1.graph.initializer.extend(model2.graph.initializer)
# input不合并
#model1.graph.input.extend(model2.graph.input)
for n in model1.graph.node:
    for i in range(len(n.input)):
        if n.input[i] == model2.graph.input[0].name:
            n.input[i] = model1.graph.input[0].name

model1.graph.output.extend(model2.graph.output)

onnx.save(model1, "merge.onnx")


#############################################runtime inference

import onnxruntime
import numpy as np

image = np.zeros((1, 3, 224, 224), dtype=np.float32)
session = onnxruntime.InferenceSession("merge.onnx", providers=["CPUExecutionProvider"])
pred = session.run(["model1-prob", "model2-prob"], {"model1-image": image})

print("Onnx Out1 = ", pred[0][0, :5])
print("Onnx Out2 = ", pred[1][0, :5])

4 onnx-parser

对于pytorch导出的onnx文件，tensorrt可以利用NvOnnxParser.h中IParser，解析onnx文件的网络到INetworkDefinition对象中。

下面是一段示例代码：

// tensorRT include
// 编译用的头文件
#include <NvInfer.h>

// onnx解析器的头文件
#include <NvOnnxParser.h>

// 推理用的运行时头文件
#include <NvInferRuntime.h>

// cuda include
#include <cuda_runtime.h>

// system include
#include <stdio.h>
#include <math.h>

#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

inline const char* severity_string(nvinfer1::ILogger::Severity t){
    switch(t){
        case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
        case nvinfer1::ILogger::Severity::kERROR:   return "error";
        case nvinfer1::ILogger::Severity::kWARNING: return "warning";
        case nvinfer1::ILogger::Severity::kINFO:    return "info";
        case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
        default: return "unknow";
    }
}

class TRTLogger : public nvinfer1::ILogger{
public:
    virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{
        if(severity <= Severity::kINFO){
            // 打印带颜色的字符，格式如下：
            // printf("\033[47;33m打印的文本\033[0m");
            // 其中 \033[ 是起始标记
            //      47    是背景颜色
            //      ;     分隔符
            //      33    文字颜色
            //      m     开始标记结束
            //      \033[0m 是终止标记
            // 其中背景颜色或者文字颜色可不写
            // 部分颜色代码 https://blog.csdn.net/ericbar/article/details/79652086
            if(severity == Severity::kWARNING){
                printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
            }
            else if(severity <= Severity::kERROR){
                printf("\033[31m%s: %s\033[0m\n", severity_string(severity), msg);
            }
            else{
                printf("%s: %s\n", severity_string(severity), msg);
            }
        }
    }
} logger;

bool build_model(){
    TRTLogger logger;

    // ----------------------------- 1. 定义 builder, config 和network -----------------------------
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);


    // ----------------------------- 2. 输入，模型结构和输出的基本信息 -----------------------------
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger);
    if(!parser->parseFromFile("demo.onnx", 1)){
        printf("Failed to parser demo.onnx\n");
        return false;
    }
    
    int maxBatchSize = 10;
    printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
    config->setMaxWorkspaceSize(1 << 28);

    // --------------------------------- 2.1 关于profile ----------------------------------
    // 如果模型有多个输入，则必须多个profile
    auto profile = builder->createOptimizationProfile();
    auto input_tensor = network->getInput(0);
    int input_channel = input_tensor->getDimensions().d[1];
    
    // 配置输入的最小、最优、最大的范围
    profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, input_channel, 3, 3));
    profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, input_channel, 3, 3));
    profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(maxBatchSize, input_channel, 5, 5));
    // 添加到配置
    config->addOptimizationProfile(profile);

    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if(engine == nullptr){
        printf("Build engine failed.\n");
        return false;
    }

    // -------------------------- 3. 序列化 ----------------------------------
    // 将模型序列化，并储存为文件
    nvinfer1::IHostMemory* model_data = engine->serialize();
    FILE* f = fopen("engine.trtmodel", "wb");
    fwrite(model_data->data(), 1, model_data->size(), f);
    fclose(f);

    // 卸载顺序按照构建顺序倒序
    model_data->destroy();
    parser->destroy();
    engine->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
    printf("Done.\n");
    return true;
}

int main(){
    build_model();
    return 0;
}

5 onnx-parser-source-code

nvonnxparser已经被nvidia给开源了， 除了使用tensorrt自带的解析器(libnvonnxparser.so)解析onnx文件，也可以使用源代码进行编译（自己进行一些定制化修改等）。下面是使用源代码加载onnx的步骤：

(参考：https://blog.csdn.net/qq_42001765/article/details/128078243)

    1. 下载onnx-tensorrt源码：https://github.com/onnx/onnx-tensorrt/releases （我这里选择8.2版本的)

      • 删除里面一些无用的文件，cmakelist，main.cpp， onnx_trt_backend.cpp等等

      • onnx-tensorrt中所有的ONNX_NAMESPACE::替换成onnx::

      • （若编译时报错，将importerContext.hpp报错的onnx::替换成ONNX_NAMESPACE::）

    1. 下载依赖项protubuf (我这里下载protobuf-3.6.1的cpp源码和编译好的protoc工具)
    1. 下载onnx-ml.proto和onnx-ml.proto文件：

    1. 利用protoc工具，生成c++解析文件，示例脚本如下：

      cd onnx
mkdir -p pbout

$protoc onnx-ml.proto --cpp_out=pbout
$protoc onnx-operators-ml.proto --cpp_out=pbout

echo Copy pbout/onnx-ml.pb.cc to ../src/onnx/onnx-ml.pb.cpp
cp pbout/onnx-ml.pb.cc           ../src/onnx/onnx-ml.pb.cpp

echo Copy pbout/onnx-operators-ml.pb.cc to ../src/onnx/onnx-operators-ml.pb.cpp
cp pbout/onnx-operators-ml.pb.cc ../src/onnx/onnx-operators-ml.pb.cpp

echo Copy pbout/onnx-ml.pb.h to ../src/onnx/onnx-ml.pb.h
cp pbout/onnx-ml.pb.h           ../src/onnx/onnx-ml.pb.h

echo Copy pbout/onnx-operators-ml.pb.h to ../src/onnx/onnx-operators-ml.pb.h
cp pbout/onnx-operators-ml.pb.h ../src/onnx/onnx-operators-ml.pb.h

echo Remove directory "pbout"
rm -rf pbout

下面是示例代码：

#include <NvInfer.h>
#include <NvInferRuntime.h>
#include <onnx-tensorrt-release-8.2-GA/NvOnnxParser.h>

#include <stdio.h>
#include <iostream>

inline const char* severity_string(nvinfer1::ILogger::Severity t) {
	switch (t) {
	case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
	case nvinfer1::ILogger::Severity::kERROR:   return "error";
	case nvinfer1::ILogger::Severity::kWARNING: return "warning";
	case nvinfer1::ILogger::Severity::kINFO:    return "info";
	case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
	default: return "unknow";
	}
}

class TRTLogger : public nvinfer1::ILogger {
public:
	virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override {
		if (severity <= Severity::kINFO) {
			// 打印带颜色的字符，格式如下：
			// printf("\033[47;33m打印的文本\033[0m");
			// 其中 \033[ 是起始标记
			//      47    是背景颜色
			//      ;     分隔符
			//      33    文字颜色
			//      m     开始标记结束
			//      \033[0m 是终止标记
			// 其中背景颜色或者文字颜色可不写
			// 部分颜色代码 https://blog.csdn.net/ericbar/article/details/79652086
			if (severity == Severity::kWARNING) {
				printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
			}
			else if (severity <= Severity::kERROR) {
				printf("\033[31m%s: %s\033[0m\n", severity_string(severity), msg);
			}
			else {
				printf("%s: %s\n", severity_string(severity), msg);
			}
		}
	}
} logger;

bool build_model() {
	TRTLogger logger;
	nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
	nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
	nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);

	nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger);
	if (!parser->parseFromFile("torch_model.onnx", 1)) {
		network->destroy();
		config->destroy();
		builder->destroy();
		printf("load onnx file failed\n");
		return false;
	}
	config->setMaxWorkspaceSize(1 << 28);
	int maxBatchSize = 10;
	auto profile = builder->createOptimizationProfile();
	auto input_tensor = network->getInput(0);
	int input_channel = input_tensor->getDimensions().d[1];
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, input_channel, 3, 3));
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, input_channel, 3, 3));
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(maxBatchSize, input_channel, 5, 5));

	config->addOptimizationProfile(profile);

	nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
	if (engine == nullptr) {
		printf("build engine failed\n");
		network->destroy();
		config->destroy();
		builder->destroy();
		return false;
	}
	nvinfer1::IHostMemory* model_data = engine->serialize();
	FILE* f = fopen("onnx_parser.trtmoddel", "wb");
	fwrite(model_data->data(), 1, model_data->size(), f);
	fclose(f);

	model_data->destroy();
	parser->destroy();
	engine->destroy();
	network->destroy();
	config->destroy();
	builder->destroy();
	printf("Done.\n");
	return true;
}

int main() {
	build_model();
	std::cin.get();
	return 0;
}
posted @ 2023-09-02 16:14  silence_cho  阅读(702)  评论(0)  编辑  收藏  举报