Add ONNX export metadata (#7353)

This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-04-09 13:27:49 +02:00
1 parent 698a5d7f26
commit 3bb233a7fb
2 files changed
+10 -1

No files matched your search

+3
View File
@@ -328,6 +328,9 @@ class DetectMultiBackend(nn.Module):
import onnxruntime
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
session = onnxruntime.InferenceSession(w, providers=providers)
meta = session.get_modelmeta().custom_metadata_map # metadata
if 'stride' in meta:
stride, names = int(meta['stride']), eval(meta['names'])
elif xml: # OpenVINO
LOGGER.info(f'Loading {w} for OpenVINO inference...')
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/