diff --git a/models/common.py b/models/common.py index d308244c4a44..2e5d5a198e33 100644 --- 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):