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来获得数据。