网站维护管理软件研发工程师

1、往onnx中写入

在训练完毕之后，我们先使用torch.onnx.export() 导出onnx模型，然后我们再使用以下代码来往metadata中写入信息：

    # Metadatad = {# 'stride': int(max(model.stride)),'names': model.names,'mean' : [0,0,0],'std' : [1,1,1],'normalize' : [1/255],'imgsz' : imgsz,'optsize' : imgsz,'min_size' : imgsz,'max_size' : imgsz,}for k, v in d.items():meta = model_onnx.metadata_props.add()meta.key, meta.value = k, str(v)onnx.save(model_onnx, f)

2、加载onnx并读取信息

model = onnx.load(model_path)# 比如我们想要读取的时类别信息
class_names = model.metadata_props[1].value

这里的model.metadata_props是一个字典列表，我们可以在debug时查看我们想要读取的信息是列表哪个索引，然后用字典的value来获得数据。