Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

DetectMultiBackend improvements (#9269)

Browse Source 
* 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>
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-09-04 01:33:38 +02:00
1 parent 5cb9fe612a
commit 63ecce60ea
1 file changed
+9 -8
Download Patch File Download Diff File Expand all files Collapse all files
models/common.py
+9 -8
View File
@@ -354,6 +354,7 @@ class DetectMultiBackend(nn.Module):
import onnxruntime import onnxruntime
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider'] providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
session = onnxruntime.InferenceSession(w, providers=providers) session = onnxruntime.InferenceSession(w, providers=providers)
output_names = [x.name for x in session.get_outputs()]
meta = session.get_modelmeta().custom_metadata_map # metadata meta = session.get_modelmeta().custom_metadata_map # metadata
if 'stride' in meta: if 'stride' in meta:
stride, names = int(meta['stride']), eval(meta['names']) stride, names = int(meta['stride']), eval(meta['names'])
@@ -372,9 +373,7 @@ class DetectMultiBackend(nn.Module):
batch_size = batch_dim.get_length() batch_size = batch_dim.get_length()
executable_network = ie.compile_model(network, device_name="CPU") # device_name="MYRIAD" for Intel NCS2 executable_network = ie.compile_model(network, device_name="CPU") # device_name="MYRIAD" for Intel NCS2
output_layer = next(iter(executable_network.outputs)) output_layer = next(iter(executable_network.outputs))
meta = Path(w).with_suffix('.yaml') stride, names = self._load_metadata(Path(w).with_suffix('.yaml')) # load metadata
if meta.exists():
stride, names = self._load_metadata(meta) # load metadata
elif engine: # TensorRT elif engine: # TensorRT
LOGGER.info(f'Loading {w} for TensorRT inference...') LOGGER.info(f'Loading {w} for TensorRT inference...')
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
@@ -476,7 +475,7 @@ class DetectMultiBackend(nn.Module):
y = self.net.forward() y = self.net.forward()
elif self.onnx: # ONNX Runtime elif self.onnx: # ONNX Runtime
im = im.cpu().numpy() # torch to numpy 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 elif self.xml: # OpenVINO
im = im.cpu().numpy() # FP32 im = im.cpu().numpy() # FP32
y = self.executable_network([im])[self.output_layer] y = self.executable_network([im])[self.output_layer]
@@ -524,7 +523,7 @@ class DetectMultiBackend(nn.Module):
y[..., :4] *= [w, h, w, h] # xywh normalized to pixels y[..., :4] *= [w, h, w, h] # xywh normalized to pixels
if isinstance(y, np.ndarray): 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 return (y, []) if val else y
def warmup(self, imgsz=(1, 3, 640, 640)): def warmup(self, imgsz=(1, 3, 640, 640)):
@@ -548,10 +547,12 @@ class DetectMultiBackend(nn.Module):
return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
@staticmethod @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 # Load metadata from meta.yaml if it exists
d = yaml_load(f) if f.exists():
return d['stride'], d['names'] # assign stride, names d = yaml_load(f)
return d['stride'], d['names'] # assign stride, names
return None, None
class AutoShape(nn.Module): class AutoShape(nn.Module):