Add ONNX export metadata (#7353)
This commit is contained in:
1 parent
698a5d7f26
commit
3bb233a7fb
2 files changed
+10
-1
No files matched your search
@@ -140,7 +140,13 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
|
||||
# Checks
|
||||
model_onnx = onnx.load(f) # load onnx model
|
||||
onnx.checker.check_model(model_onnx) # check onnx model
|
||||
# LOGGER.info(onnx.helper.printable_graph(model_onnx.graph)) # print
|
||||
|
||||
# Metadata
|
||||
d = {'stride': int(max(model.stride)), 'names': model.names}
|
||||
for k, v in d.items():
|
||||
meta = model_onnx.metadata_props.add()
|
||||
meta.key, meta.value = k, str(v)
|
||||
onnx.save(model_onnx, f)
|
||||
|
||||
# Simplify
|
||||
if simplify:
|
||||
|
||||
Reference in new issue
Block a user