DetectMultiBackend improvements (ultralytics#9269)

* Update common.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Update common.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Update common.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Update common.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
ctjanuhowski · Sep 8, 2022 · 59b9ad1 · 59b9ad1
1 parent 65bc4ba
commit 59b9ad1
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/models/common.py b/models/common.py
@@ -354,6 +354,7 @@ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False,
             import onnxruntime
             providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
             session = onnxruntime.InferenceSession(w, providers=providers)
+            output_names = [x.name for x in session.get_outputs()]
             meta = session.get_modelmeta().custom_metadata_map  # metadata
             if 'stride' in meta:
                 stride, names = int(meta['stride']), eval(meta['names'])
@@ -372,9 +373,7 @@ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False,
                 batch_size = batch_dim.get_length()
             executable_network = ie.compile_model(network, device_name="CPU")  # device_name="MYRIAD" for Intel NCS2
             output_layer = next(iter(executable_network.outputs))
-            meta = Path(w).with_suffix('.yaml')
-            if meta.exists():
-                stride, names = self._load_metadata(meta)  # load metadata
+            stride, names = self._load_metadata(Path(w).with_suffix('.yaml'))  # load metadata
         elif engine:  # TensorRT
             LOGGER.info(f'Loading {w} for TensorRT inference...')
             import tensorrt as trt  # https://developer.nvidia.com/nvidia-tensorrt-download
@@ -476,7 +475,7 @@ def forward(self, im, augment=False, visualize=False, val=False):
             y = self.net.forward()
         elif self.onnx:  # ONNX Runtime
             im = im.cpu().numpy()  # torch to numpy
-            y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
+            y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})[0]
         elif self.xml:  # OpenVINO
             im = im.cpu().numpy()  # FP32
             y = self.executable_network([im])[self.output_layer]
@@ -524,7 +523,7 @@ def forward(self, im, augment=False, visualize=False, val=False):
             y[..., :4] *= [w, h, w, h]  # xywh normalized to pixels
 
         if isinstance(y, np.ndarray):
-            y = torch.tensor(y, device=self.device)
+            y = torch.from_numpy(y).to(self.device)
         return (y, []) if val else y
 
     def warmup(self, imgsz=(1, 3, 640, 640)):
@@ -548,10 +547,12 @@ def _model_type(p='path/to/model.pt'):
         return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
 
     @staticmethod
-    def _load_metadata(f='path/to/meta.yaml'):
+    def _load_metadata(f=Path('path/to/meta.yaml')):
         # Load metadata from meta.yaml if it exists
-        d = yaml_load(f)
-        return d['stride'], d['names']  # assign stride, names
+        if f.exists():
+            d = yaml_load(f)
+            return d['stride'], d['names']  # assign stride, names
+        return None, None
 
 
 class AutoShape(nn.Module):