【深度学习|onnx】往onnx中写入训练的超参或者类别等信息，并在推理时读取
1、往onnx中写入
在训练完毕之后，我们先使用torch.onnx.export() 导出onnx模型，然后我们再使用以下代码来往metadata中写入信息：
# Metadata
d = {
# '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来获得数据。