ONNX网络模型解析

        ONNX是一种针对机器学习所设计的开放式的文件格式,用于存储训练好的模型。它使得不同的人工智能框架(如Pytorch, MXNet)可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软,亚马逊 ,Facebook 和 IBM 等公司共同开发,以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有: Caffe2, PyTorch, MXNet,ML.NET,TensorRT 和 Microsoft CNTK,并且 TensorFlow 也非官方的支持ONNX。(参考维基百科)

        如需自定义网络模型结构,尝试使用自定义的结构保存网络模型,首先需要解析ONNX的net结构。ONNX使用google定义的报文格式protocol buffer,用于RPC 系统和持续数据存储系统。

        Protocol Buffers 是一种轻便高效的结构化数据存储格式,可以用于结构化数据串行化,或者说序列化。它很适合做数据存储或 RPC 数据交换格式。可用于通讯协议、数据存储等领域的语言无关、平台无关、可扩展的序列化结构数据格式。目前提供了 C++、Java、Python 三种语言的 API,本文使用了C++接口解析ONNX模型。

        Protocol buffer结构通常会定义要给proto文件,通过ONNXgithub链接下载onnx.protohttps://github.com/onnx/onnx。解析模型用到的结构主要如下:

  1. ModelProto:最高级别的结构,定义了整个网络模型结构;
  2. GraphProto: graph定义了模型的计算逻辑以及带有参数的node节点,组成一个有向图结构;
  3. NodeProto: 网络有向图的各个节点OP的结构,通常称为层,例如conv,relu层;
  4. AttributeProto:各OP的参数,通过该结构访问,例如:conv层的stride,dilation等;
  5. TensorProto: 序列化的tensor value,一般weight,bias等常量均保存为该种结构;
  6. TensorShapeProto:网络的输入shape以及constant输入tensor的维度信息均保存为该种结构;
  7. TypeProto:表示ONNX数据类型。

        具体解析流程是读取.onnx文件,获得一个model结构,通过model结构访问到graph结构,然后通过graph访问整个网络的所有node以及input,output,通过node结构可以访问到OP的参数。

下面给出解析demo:

void ReadProtoFromBinaryFile(const char* filename, google::protobuf::Message* proto)
 {
	int fd = open(filename, O_RDONLY);
	google::protobuf::io::FileInputStream* raw_input = new google::protobuf::io::FileInputStream(fd);
	google::protobuf::io::CodedInputStream* coded_input = new google::protobuf::io::CodedInputStream(raw_input);
	coded_input->SetTotalBytesLimit(INT_MAX, 536870912);
	bool success = proto->ParseFromCodedStream(coded_input);
	//bool success = proto->ParseFromZeroCopyStream(raw_input);
	delete coded_input;
	delete raw_input;
	close(fd);
	if (success != true)
	{
		exit(1);
	}
}
int main()
{
	char* ch = "inception_v3.onnx";
	onnx::ModelProto model_data;
	ReadProtoFromBinaryFile(ch, &model_data);    //读取文件保存为modelproto结构
	onnx::GraphProto graph = model_data.graph();  //访问graph结构
	int num = graph.node_size();                    //node节点个数
	int input_size = graph.input_size();            //网络输入个数,input以及各层常量输入
	for (int i = 0; i < input_size; ++i)
	{
		const std::string name = graph.input(i).name();
		onnx::TypeProto type = graph.input(i).type();
		onnx::TensorShapeProto shape = type.tensor_type().shape();//输入维度
		for (int i = 0; i < shape.dim_size(); i++)
		{
			std::cout << shape.dim(i).dim_value();
			std::cout<< std::endl;
		}
	}
	int output_size = graph.output_size();    //网络output个数
	for (int i = 0; i < num; i++)    //遍历每个node结构
	{
		const onnx::NodeProto node = graph.node(i);
		std::string node_name = node.name();
		std::cout <<"cur node name:"<< node_name << std::endl;
		const ::google::protobuf::RepeatedPtrField< ::onnx::AttributeProto> attr = node.attribute();        //每个node结构的参数信息
		const std::string type = node.op_type();
		int in_size = node.input_size();	
		int out_size = node.output_size();
	}
	return 0;
}

 

  • 3
    点赞
  • 27
    收藏
    觉得还不错? 一键收藏
  • 2
    评论
要贯通 YOLOv7 网络模型并导出 ONNX 模型,您可以按照以下步骤进行: 1. 下载 YOLOv7 模型的代码库,例如 https://github.com/WongKinYiu/yolov7。 2. 在本地或云端环境中安装 PyTorch 和其他必备的 Python 包。 3. 在数据集上训练 YOLOv7 模型,或者下载预训练模型。 4. 修改代码库中的 export.py 文件,以便将模型导出为 ONNX 格式。您需要指定输入和输出张量的名称和形状,并将模型保存为 ONNX 文件。 5. 运行 export.py 文件并检查导出的 ONNX 模型是否可用。 以下是一个简单的示例代码,展示了如何导出 YOLOv7 模型ONNX 格式: ``` python import torch import torch.nn as nn import argparse # 导入yolov7模型 from models.yolov7 import YOLOv7 parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='path/to/weights.pt', help='path to weights file') parser.add_argument('--img_size', type=int, default=416, help='size of each image dimension') parser.add_argument('--batch_size', type=int, default=1, help='size of the batches') parser.add_argument('--output', type=str, default='path/to/output.onnx', help='output ONNX file path') opt = parser.parse_args() # 加载yolov7模型 model = YOLOv7(opt.img_size, num_classes=80) model.load_state_dict(torch.load(opt.weights, map_location='cpu')['model']) # 设置输入张量的名称和形状 input_names = ['input_0'] input_shapes = [(opt.batch_size, 3, opt.img_size, opt.img_size)] # 设置输出张量的名称和形状 output_names = ['output_0', 'output_1', 'output_2'] output_shapes = [(opt.batch_size, 255, opt.img_size//32, opt.img_size//32), (opt.batch_size, 255, opt.img_size//16, opt.img_size//16), (opt.batch_size, 255, opt.img_size//8, opt.img_size//8)] # 将模型转换为ONNX格式 torch.onnx.export(model, torch.zeros(*input_shapes), opt.output, input_names=input_names, output_names=output_names, output_shapes=output_shapes, opset_version=11) print(f'Exported ONNX model to {opt.output}') ``` 这里我们假设您已经按照上述步骤安装了 PyTorch 和其他必备的 Python 包。在这个示例中,我们使用 argparse 模块来解析命令行参数,指定权重文件的路径、图像大小、批处理大小和输出 ONNX 文件的路径。然后,我们加载 YOLOv7 模型,并指定输入和输出张量的名称和形状。最后,我们使用 `torch.onnx.export()` 函数将模型保存为 ONNX 文件。 请注意,导出模型时需要指定适当的 opset 版本。在这个示例中,我们使用的是 opset_version=11,这是与 PyTorch 1.7.x 版本兼容的最新版本。 希望这个简单的示例能够帮助您了解如何贯通 YOLOv7 网络模型并导出 ONNX 模型

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值