1. 创建一个pytorch模型
这里我用的U2Net,直接加载好训练出的权重
model = U2Net(class_nums=4) model.load_state_dict(torch.load(checkpoint_path))
2. 将pytorch模型转成onnx格式
x = torcg.randn(1,3,512,512) with torch.no_grad(): torch.onnx.export( model, x, "onnx_file_name.onnx", opset_version=11, input_names=['input'], output_names=['output'] )
torch.onnx.export()函数用于将pytorch模型转成onnx格式的函数,其中的参数如下
- model:需要转换的模型,加载好权重
- args:模型的任意一组输入,注意维度
- f:导出的onnx模型的文件名
- opset_version:表示ONNX算子集的版本,随着深度学习的发展,新算子会不断诞生,算子集就需要扩充,截至20230824,官网已经发布20个版本
- input_names:输入tensor的名称
- output_names:输出tensor的名称
转换成功会在指定的路径f下生成onnx模型文件
可用onnx中的函数验证模型文件是否正确
import onnx onnx_model = onnx.load("onnx_file_name.onnx") try: onnx.checker.check_model(onnx_model) except Exception as ex: print(f'ERROR: {ex}') else: print(f'model correct')
可以使用netron查看转换的onnx模型的详细内容 链接 https://netron.app/
点击某个算子节点,可以查看算子的具体信息,每个算子记录了算子的属性、图结构和权重三类信息
- 算子属性信息(attributes),对于卷积来说,算子属性包含了卷积核的大小、步长等属性
- 图结构信息(node properties),算子节点在计算图中的名称、邻边的信息,如上图节点名为conv2,输入数据叫683,权重叫1324
- 权重信息指算子存储的网络权重信息,点击w后的+号可以查看权重信息的具体内容
3. 推理引擎ONNXRuntime
使用onnxruntime可以运行输出的onnx模型,代码如下
input_img 是一个与onnx模型输入维度一致的归一化的np.array model = onnxruntime.InferenceSession(onnx_file) inputs = {'input':input_img} output = model.run(['output'],inputs)[0]
onnxruntime.InferenceSession用于获取一个ONNX Runtime的推理器,其参数为用于推理的ONNX模型文件
run方法用于模型推理,第一个参数为输出张量名的列表,第二个参数为输入值的字典,key为导出模型时设置的张量名,value为输入张量
标签:output,onnx,模型,pytorch,算子,input,model From: https://www.cnblogs.com/Liang-ml/p/17653716.html