onnx模型修改动态batch size和增加标签数据

最新推荐文章于 2023-12-19 13:41:10 发布

jstzwjr

最新推荐文章于 2023-12-19 13:41:10 发布

阅读量558

点赞数

文章标签：深度学习

本文链接：https://blog.csdn.net/qq_17127427/article/details/130969415

版权

该代码段展示了如何利用onnx库和onnx-simplifier工具对名为ace2p的ONNX模型进行简化。首先导入必要的库，定义类别标签，然后读取并简化模型。接着，添加元数据到模型，并修改输入形状。最后，保存简化后的模型到指定路径。

摘要由CSDN通过智能技术生成

import onnx
from onnxsim import simplify

labels = ['Background', 'Hat', 'Hair', 'Glove', 'Sunglasses', 'Upper-clothes', 'Dress', 'Coat',
          'Socks', 'Pants', 'Jumpsuits', 'Scarf', 'Skirt', 'Face', 'Left-arm', 'Right-arm',
          'Left-leg', 'Right-leg', 'Left-shoe', 'Right-shoe']
names = {cls_id: label for cls_id, label in enumerate(labels)}

model_name = "ace2p"

sim_model, ret = simplify(f'{model_name}.onnx')

meta = sim_model.metadata_props.add()
meta.key, meta.value = "names", str(names)

graph = sim_model.graph
graph.input[0].type.tensor_type.shape.dim[0].dim_param = 'None'
onnx.save(sim_model, f'/wjr/weights/{model_name}.onnx')
···