From 8c3f2336ddb84e659edcfd0135d5d83135329442 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 12:10:44 -0800 Subject: [PATCH 01/25] Add benchmarks.py --- utils/benchmarks.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 utils/benchmarks.py diff --git a/utils/benchmarks.py b/utils/benchmarks.py new file mode 100644 index 000000000000..e69de29bb2d1 From 9a52627aa22811da3c28a4faa164342285c0d19d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 12:26:00 -0800 Subject: [PATCH 02/25] Update --- utils/benchmarks.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index e69de29bb2d1..2d58a88fdc70 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -0,0 +1,39 @@ +# YOLOv5 🚀 by Ultralytics, GPL-3.0 license +""" +Run YOLOv5 benchmarks on all supported export formats + +Format | `export.py --include` | Model +--- | --- | --- +PyTorch | - | yolov5s.pt +TorchScript | `torchscript` | yolov5s.torchscript +ONNX | `onnx` | yolov5s.onnx +OpenVINO | `openvino` | yolov5s_openvino_model/ +TensorRT | `engine` | yolov5s.engine +CoreML | `coreml` | yolov5s.mlmodel +TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/ +TensorFlow GraphDef | `pb` | yolov5s.pb +TensorFlow Lite | `tflite` | yolov5s.tflite +TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite +TensorFlow.js | `tfjs` | yolov5s_web_model/ + +Usage: + $ python utils/benchmarks.py --weights yolov5s.pt --img 640 +""" + +import sys +from pathlib import Path + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[1] # YOLOv5 root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +# ROOT = ROOT.relative_to(Path.cwd()) # relative + +import export, val + +suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', + '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] + +export.run() + +val.run() From d9933ad26904ab8e9c36024d8acbdcda27d55836 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 18:16:31 -0500 Subject: [PATCH 03/25] Add requirements --- export.py | 4 ++++ utils/benchmarks.py | 41 +++++++++++++++++++++++++++++++++++++---- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/export.py b/export.py index 589b381e035a..28e13358fd64 100644 --- a/export.py +++ b/export.py @@ -16,6 +16,10 @@ TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite TensorFlow.js | `tfjs` | yolov5s_web_model/ +Requirements: + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU + Usage: $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ... diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 2d58a88fdc70..75872992a2f8 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -16,10 +16,15 @@ TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite TensorFlow.js | `tfjs` | yolov5s_web_model/ +Requirements: + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU + Usage: $ python utils/benchmarks.py --weights yolov5s.pt --img 640 """ +import argparse import sys from pathlib import Path @@ -30,10 +35,38 @@ # ROOT = ROOT.relative_to(Path.cwd()) # relative import export, val +from utils.general import print_args + + +def run(weights=ROOT / 'yolov5s.pt', # weights path + imgsz=(640, 640), # inference size h,w + batch_size=1, # batch size + ): + formats = 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' + suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', + 'edgetpu.tflite', '_web_model'] + + for format in formats: + files = export.run(weights=weights, imgsz=imgsz, include=[format]) + result = val.run(weights=files[-1], imgsz=imgsz) + print(result) + + +def parse_opt(): + parser = argparse.ArgumentParser() + parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') + parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') + parser.add_argument('--batch-size', type=int, default=1, help='batch size') + opt = parser.parse_args() + opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand + print_args(FILE.stem, opt) + return opt + -suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', - '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] +def main(opt): + run(**vars(opt)) -export.run() -val.run() +if __name__ == "__main__": + opt = parse_opt() + main(opt) From 818197c06dcdccfaccc7e25924480ffdebbaffb0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 18:40:52 -0500 Subject: [PATCH 04/25] Updates --- utils/benchmarks.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 6af01fd7d8cd..3f7c6dee0e94 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -39,7 +39,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path - imgsz=(640, 640), # inference size h,w + imgsz=640, # inference size (pixels) batch_size=1, # batch size data=ROOT / 'data/coco128.yaml', # dataset.yaml path ): @@ -47,20 +47,26 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] - for format in formats: - files = export.run(weights=weights, imgsz=imgsz, include=[format]) - result = val.run(data=data, weights=files[-1], imgsz=imgsz) - print(result) + y = [] + for f in formats[:2]: + file = export.run(weights=weights, imgsz=[imgsz], include=[f])[-1] + result = val.run(data=data, weights=file, imgsz=imgsz, batch_size=batch_size, plots=False) + m = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) + t = result[2] # times (preprocess, inference, postprocess) + + # y.append([f, file, *m, *t]) + y.append([f, Path(file).name, m[3], t[1]]) # mAP, t_inference + + print(y) def parse_opt(): parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') - parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') + parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=64, help='inference size (pixels)') parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') opt = parser.parse_args() - opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand print_args(FILE.stem, opt) return opt From 93bd69ee62fcd4dd9a5a1dd2c3cfab2d9591b02a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 18:45:58 -0500 Subject: [PATCH 05/25] Updates --- utils/benchmarks.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 3f7c6dee0e94..63f04a75416f 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -43,20 +43,16 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path batch_size=1, # batch size data=ROOT / 'data/coco128.yaml', # dataset.yaml path ): - formats = 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' - suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', - 'edgetpu.tflite', '_web_model'] + formats = 'openvino', 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' + # suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] y = [] - for f in formats[:2]: - file = export.run(weights=weights, imgsz=[imgsz], include=[f])[-1] + for f in formats[:3]: + file = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f])[-1] result = val.run(data=data, weights=file, imgsz=imgsz, batch_size=batch_size, plots=False) m = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) t = result[2] # times (preprocess, inference, postprocess) - - # y.append([f, file, *m, *t]) y.append([f, Path(file).name, m[3], t[1]]) # mAP, t_inference - print(y) From 83b7e077014a0d1313a2d4323f6e3833c9524d08 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 19:06:04 -0500 Subject: [PATCH 06/25] Updates --- utils/benchmarks.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 63f04a75416f..13cc176a99b5 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -28,14 +28,17 @@ import sys from pathlib import Path +import pandas as pd + FILE = Path(__file__).resolve() ROOT = FILE.parents[1] # YOLOv5 root directory if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH # ROOT = ROOT.relative_to(Path.cwd()) # relative + import export, val -from utils.general import print_args +from utils.general import LOGGER, print_args def run(weights=ROOT / 'yolov5s.pt', # weights path @@ -43,7 +46,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path batch_size=1, # batch size data=ROOT / 'data/coco128.yaml', # dataset.yaml path ): - formats = 'openvino', 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' + formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' # suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] y = [] @@ -53,7 +56,10 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path m = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) t = result[2] # times (preprocess, inference, postprocess) y.append([f, Path(file).name, m[3], t[1]]) # mAP, t_inference - print(y) + + py = pd.DataFrame(y, columns=['Format', 'Weights', 'mAP@0.5:0.95', 'Inference time (ms)']) + LOGGER.info('\nBenchmarks finished.') + LOGGER.info(py) def parse_opt(): From aeb6ee448d5f230c948e45d156da6f29dfe60f49 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 19:07:02 -0500 Subject: [PATCH 07/25] Updates --- utils/benchmarks.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 13cc176a99b5..f54f9e5ecb56 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -36,7 +36,6 @@ sys.path.append(str(ROOT)) # add ROOT to PATH # ROOT = ROOT.relative_to(Path.cwd()) # relative - import export, val from utils.general import LOGGER, print_args @@ -50,7 +49,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path # suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] y = [] - for f in formats[:3]: + for f in formats[:-1]: file = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f])[-1] result = val.run(data=data, weights=file, imgsz=imgsz, batch_size=batch_size, plots=False) m = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) From 89dbad3ad2b02a9dcc9fae338965bfe6df9d22d2 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 20:01:53 -0500 Subject: [PATCH 08/25] Updates --- models/common.py | 16 ++++++++++++++++ utils/benchmarks.py | 9 ++++----- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/models/common.py b/models/common.py index 346fa37ae2d0..dc71d7863885 100644 --- a/models/common.py +++ b/models/common.py @@ -35,6 +35,22 @@ def autopad(k, p=None): # kernel, padding return p +def model_type(p='path/to/model.pt'): + # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx + formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', \ + 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' + suffixes = '.pt', '.torchscript', '.onnx', '.xml', '_openvino_model', '.engine', '.mlmodel', \ + '_saved_model', '.pb', '.tflite', '_edgetpu.tflite', '_web_model' + + # suffix = Path(p).suffix.lower() + # suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.tflite', '.pb', '', '.mlmodel'] + # check_suffix(p, suffixes) # check weights have acceptable suffix + pt, jit, onnx, xml, xml2, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = (p.endswith(x) for x in suffixes) + tflite = tflite and not edgetpu + xml |= xml2 # *.xml or *_openvino_model + return pt, jit, onnx, xml, xml2, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs + + class Conv(nn.Module): # Standard convolution def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups diff --git a/utils/benchmarks.py b/utils/benchmarks.py index f54f9e5ecb56..5be78d511286 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -45,13 +45,12 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path batch_size=1, # batch size data=ROOT / 'data/coco128.yaml', # dataset.yaml path ): - formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' - # suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '_saved_model', '.pb', '.tflite', 'edgetpu.tflite', '_web_model'] y = [] + formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' for f in formats[:-1]: - file = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f])[-1] - result = val.run(data=data, weights=file, imgsz=imgsz, batch_size=batch_size, plots=False) + file = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] + result = val.run(data=data, weights=file, imgsz=imgsz, batch_size=batch_size, plots=False, device='cpu') m = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) t = result[2] # times (preprocess, inference, postprocess) y.append([f, Path(file).name, m[3], t[1]]) # mAP, t_inference @@ -64,7 +63,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path def parse_opt(): parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') - parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=64, help='inference size (pixels)') + parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') opt = parser.parse_args() From 7dfd503feadfafb6e6e9aca89f48c075de9aec0c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 25 Jan 2022 23:27:14 -0500 Subject: [PATCH 09/25] Updates --- models/common.py | 35 +++++++++++++++-------------------- utils/benchmarks.py | 30 ++++++++++++++++++------------ 2 files changed, 33 insertions(+), 32 deletions(-) diff --git a/models/common.py b/models/common.py index dc71d7863885..3debcfd54e43 100644 --- a/models/common.py +++ b/models/common.py @@ -35,22 +35,6 @@ def autopad(k, p=None): # kernel, padding return p -def model_type(p='path/to/model.pt'): - # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx - formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', \ - 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' - suffixes = '.pt', '.torchscript', '.onnx', '.xml', '_openvino_model', '.engine', '.mlmodel', \ - '_saved_model', '.pb', '.tflite', '_edgetpu.tflite', '_web_model' - - # suffix = Path(p).suffix.lower() - # suffixes = ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.tflite', '.pb', '', '.mlmodel'] - # check_suffix(p, suffixes) # check weights have acceptable suffix - pt, jit, onnx, xml, xml2, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = (p.endswith(x) for x in suffixes) - tflite = tflite and not edgetpu - xml |= xml2 # *.xml or *_openvino_model - return pt, jit, onnx, xml, xml2, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs - - class Conv(nn.Module): # Standard convolution def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups @@ -310,10 +294,7 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None): super().__init__() w = str(weights[0] if isinstance(weights, list) else weights) - suffix = Path(w).suffix.lower() - suffixes = ['.pt', '.torchscript', '.onnx', '.engine', '.tflite', '.pb', '', '.mlmodel', '.xml'] - check_suffix(w, suffixes) # check weights have acceptable suffix - pt, jit, onnx, engine, tflite, pb, saved_model, coreml, xml = (suffix == x for x in suffixes) # backends + pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults w = attempt_download(w) # download if not local if data: # data.yaml path (optional) @@ -348,6 +329,8 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None): check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/ import openvino.inference_engine as ie core = ie.IECore() + if not Path(w).is_file(): # if not *.xml + w = next(Path(w).glob('*.xml')) # get *.xml file from *_openvino_model dir network = core.read_network(model=w, weights=Path(w).with_suffix('.bin')) # *.xml, *.bin paths executable_network = core.load_network(network, device_name='CPU', num_requests=1) elif engine: # TensorRT @@ -477,6 +460,18 @@ def warmup(self, imgsz=(1, 3, 640, 640), half=False): im = torch.zeros(*imgsz).to(self.device).type(torch.half if half else torch.float) # input image self.forward(im) # warmup + @staticmethod + def model_type(p='path/to/model.pt'): + # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx + check_suffix(p, ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '.pb', '.tflite', '']) # checks + ends = '.pt', '.torchscript', '.onnx', '.xml', '_openvino_model', '.engine', '.mlmodel', \ + '_saved_model', '.pb', '.tflite', '_edgetpu.tflite', '_web_model' + p = Path(p).name # eliminate trailing separators + pt, jit, onnx, xml, xml2, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = (p.endswith(x) for x in ends) + xml |= xml2 # *.xml or *_openvino_model + tflite &= not edgetpu # *.tflite + return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs + class AutoShape(nn.Module): # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 5be78d511286..557718f5b3ec 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -26,6 +26,7 @@ import argparse import sys +import time from pathlib import Path import pandas as pd @@ -45,19 +46,24 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path batch_size=1, # batch size data=ROOT / 'data/coco128.yaml', # dataset.yaml path ): - - y = [] - formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' - for f in formats[:-1]: - file = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] - result = val.run(data=data, weights=file, imgsz=imgsz, batch_size=batch_size, plots=False, device='cpu') - m = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) - t = result[2] # times (preprocess, inference, postprocess) - y.append([f, Path(file).name, m[3], t[1]]) # mAP, t_inference - - py = pd.DataFrame(y, columns=['Format', 'Weights', 'mAP@0.5:0.95', 'Inference time (ms)']) - LOGGER.info('\nBenchmarks finished.') + y, t = [], time.time() + formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', \ + 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' + for f in formats[:4]: + try: + w = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] + result = val.run(data=data, weights=w, imgsz=imgsz, batch_size=batch_size, plots=False, device='cpu') + metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) + speeds = result[2] # times (preprocess, inference, postprocess) + y.append([Path(w).name, metrics[3], speeds[1]]) # mAP, t_inference + except Exception as e: + LOGGER.warning(f'WARNING: Benchmark failure for {f}: {e}') + y.append([f, None, None]) # mAP, t_inference + + py = pd.DataFrame(y, columns=['Weights', 'mAP@0.5:0.95', 'Inference time (ms)']) + LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)') LOGGER.info(py) + return py def parse_opt(): From 7bb0612898b439ec704ea8ba1ce26fe6c0bf10f1 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 26 Jan 2022 14:36:18 +0100 Subject: [PATCH 10/25] dataset autodownload from root --- utils/general.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/utils/general.py b/utils/general.py index e9f5ec2ac128..19fd347971cf 100755 --- a/utils/general.py +++ b/utils/general.py @@ -385,8 +385,8 @@ def check_dataset(data, autodownload=True): # Download (optional) extract_dir = '' if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip - download(data, dir='../datasets', unzip=True, delete=False, curl=False, threads=1) - data = next((Path('../datasets') / Path(data).stem).rglob('*.yaml')) + download(data, dir=str(ROOT / '../datasets'), unzip=True, delete=False, curl=False, threads=1) + data = next((ROOT / '../datasets' / Path(data).stem).rglob('*.yaml')) extract_dir, autodownload = data.parent, False # Read yaml (optional) @@ -405,7 +405,7 @@ def check_dataset(data, autodownload=True): data['names'] = [f'class{i}' for i in range(data['nc'])] # assign class names if missing train, val, test, s = (data.get(x) for x in ('train', 'val', 'test', 'download')) if val: - val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path + val = [Path(ROOT / x).resolve() for x in (val if isinstance(val, list) else [val])] # val path if not all(x.exists() for x in val): LOGGER.info('\nDataset not found, missing paths: %s' % [str(x) for x in val if not x.exists()]) if s and autodownload: # download script From 0cfd8eac84dea7f353c0cc4f56d24cc9c9223e2c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 3 Feb 2022 20:32:37 +0100 Subject: [PATCH 11/25] Update --- export.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/export.py b/export.py index 09c50baa415a..87361ceef95d 100644 --- a/export.py +++ b/export.py @@ -320,11 +320,14 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}' if subprocess.run(cmd, shell=True).returncode != 0: LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}') + sudo = subprocess.run('sudo --help', shell=True).returncode == 0 # sudo installed on system for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -', 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', 'sudo apt-get update', 'sudo apt-get install edgetpu-compiler']: - subprocess.run(c, shell=True, check=True) + c_ = c if sudo else c.replace('sudo ', '') + print(c_) + subprocess.run(c_, shell=True, check=True) ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...') From 7cf5f6526edfd4589948eb8d5ab77bda975c799b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 3 Feb 2022 20:36:40 +0100 Subject: [PATCH 12/25] Redirect to /dev/null --- export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index 87361ceef95d..0f891420b198 100644 --- a/export.py +++ b/export.py @@ -315,12 +315,12 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ try: - cmd = 'edgetpu_compiler --version' + cmd = 'edgetpu_compiler --version >/dev/null' help_url = 'https://coral.ai/docs/edgetpu/compiler/' assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}' if subprocess.run(cmd, shell=True).returncode != 0: LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}') - sudo = subprocess.run('sudo --help', shell=True).returncode == 0 # sudo installed on system + sudo = subprocess.run('sudo --help >/dev/null', shell=True).returncode == 0 # sudo installed on system for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -', 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', 'sudo apt-get update', From c3c575e2dae2c4931a9c9a0574101139e80b51d5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 3 Feb 2022 20:38:31 +0100 Subject: [PATCH 13/25] sudo --help --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index 0f891420b198..d001edbc1a54 100644 --- a/export.py +++ b/export.py @@ -320,7 +320,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}' if subprocess.run(cmd, shell=True).returncode != 0: LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}') - sudo = subprocess.run('sudo --help >/dev/null', shell=True).returncode == 0 # sudo installed on system + sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -', 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', 'sudo apt-get update', From 7318376bc6d169766e1de0cab0babe907362a349 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 3 Feb 2022 20:39:16 +0100 Subject: [PATCH 14/25] Cleanup --- export.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/export.py b/export.py index d001edbc1a54..6fcab69139d7 100644 --- a/export.py +++ b/export.py @@ -325,9 +325,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', 'sudo apt-get update', 'sudo apt-get install edgetpu-compiler']: - c_ = c if sudo else c.replace('sudo ', '') - print(c_) - subprocess.run(c_, shell=True, check=True) + subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True) ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...') From a7fbd4a480aadafbb0b427a0073b5435cfcf375c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 14:22:28 +0100 Subject: [PATCH 15/25] Add exports pd df --- export.py | 17 +++++++++++++++++ utils/benchmarks.py | 8 ++++---- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/export.py b/export.py index 444ab57f5c96..37670b428f44 100644 --- a/export.py +++ b/export.py @@ -52,6 +52,7 @@ import warnings from pathlib import Path +import pandas as pd import torch import torch.nn as nn from torch.utils.mobile_optimizer import optimize_for_mobile @@ -72,6 +73,22 @@ from utils.torch_utils import select_device +def export_formats(): + # YOLOv5 export formats + x = [['PyTorch', '-', '.pt'], + ['TorchScript', 'torchscript', '.torchscript'], + ['ONNX', 'onnx', '.onnx'], + ['OpenVINO', 'openvino', '_openvino_model/'], + ['TensorRT', 'engine', '.engine'], + ['CoreML', 'coreml', '.mlmodel'], + ['TensorFlow SavedModel', 'saved_model', '_saved_model/'], + ['TensorFlow GraphDef', 'pb', '.pb'], + ['TensorFlow Lite', 'tflite', '.tflite'], + ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite'], + ['TensorFlow.js', 'tfjs', '_web_model/']] + return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix']) + + def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')): # YOLOv5 TorchScript model export try: diff --git a/utils/benchmarks.py b/utils/benchmarks.py index b702c91b936f..0212c782221e 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -47,11 +47,11 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path data=ROOT / 'data/coco128.yaml', # dataset.yaml path ): y, t = [], time.time() - formats = 'torch', 'torchscript', 'onnx', 'openvino', 'engine', 'coreml', \ - 'saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs' - for f in formats[:]: + formats = export.export_formats() + for i, (name, f, suffix) in formats.iterrows(): # index, (name, file, suffix) try: - w = weights if f == 'torch' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] + w = weights if f == '-' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] + assert str(w).endswith(suffix), 'export failed' result = val.run(data=data, weights=w, imgsz=imgsz, batch_size=batch_size, plots=False, device='cpu') metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) speeds = result[2] # times (preprocess, inference, postprocess) From 8097eccbcdb2f49faaacdb2f7c880c8844fd0b94 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 15:39:20 +0100 Subject: [PATCH 16/25] Updates --- utils/__init__.py | 2 +- utils/benchmarks.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/utils/__init__.py b/utils/__init__.py index 4658ed6473cd..2f2eb1dab8e7 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -11,7 +11,7 @@ def notebook_init(verbose=True): import os import shutil - from utils.general import check_requirements, emojis, is_colab + from utils.general import LOGGER, check_requirements, emojis, is_colab from utils.torch_utils import select_device # imports check_requirements(('psutil', 'IPython')) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 0212c782221e..176c2183e579 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -38,6 +38,7 @@ # ROOT = ROOT.relative_to(Path.cwd()) # relative import export, val +from utils import notebook_init from utils.general import LOGGER, print_args @@ -55,12 +56,15 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path result = val.run(data=data, weights=w, imgsz=imgsz, batch_size=batch_size, plots=False, device='cpu') metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) speeds = result[2] # times (preprocess, inference, postprocess) - y.append([Path(w).name, metrics[3], speeds[1]]) # mAP, t_inference + y.append([name, metrics[3], speeds[1]]) # mAP, t_inference except Exception as e: - LOGGER.warning(f'WARNING: Benchmark failure for {f}: {e}') - y.append([f, None, None]) # mAP, t_inference + LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}') + y.append([name, None, None]) # mAP, t_inference - py = pd.DataFrame(y, columns=['Weights', 'mAP@0.5:0.95', 'Inference time (ms)']) + LOGGER.info('\n') + parse_opt() + notebook_init() # print system info + py = pd.DataFrame(y, columns=['Format', 'mAP@0.5:0.95', 'Inference time (ms)']) LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)') LOGGER.info(py) return py From d5f8eb3213b51da022fe00403d9927bf7c339a1a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 16:18:24 +0100 Subject: [PATCH 17/25] Updates --- utils/benchmarks.py | 2 +- val.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 176c2183e579..10fa807282b6 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -53,7 +53,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path try: w = weights if f == '-' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] assert str(w).endswith(suffix), 'export failed' - result = val.run(data=data, weights=w, imgsz=imgsz, batch_size=batch_size, plots=False, device='cpu') + result = val.run(data, w, batch_size, imgsz=imgsz, plots=False, device='cpu', task='benchmark') metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) speeds = result[2] # times (preprocess, inference, postprocess) y.append([name, metrics[3], speeds[1]]) # mAP, t_inference diff --git a/val.py b/val.py index 90debaf0dd60..78abbda8231a 100644 --- a/val.py +++ b/val.py @@ -163,9 +163,10 @@ def run(data, # Dataloader if not training: model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz), half=half) # warmup - pad = 0.0 if task == 'speed' else 0.5 + pad = 0.0 if task in ('speed', 'benchmark') else 0.5 + rect = False if task == 'benchmark' else pt # square inference for benchmarks task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images - dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=pt, + dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=rect, workers=workers, prefix=colorstr(f'{task}: '))[0] seen = 0 From 257186c992df162ce0645aa63bcda2bc234b7675 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 16:20:32 +0100 Subject: [PATCH 18/25] Updates --- utils/benchmarks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 10fa807282b6..2e407cc93158 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -66,7 +66,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path notebook_init() # print system info py = pd.DataFrame(y, columns=['Format', 'mAP@0.5:0.95', 'Inference time (ms)']) LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)') - LOGGER.info(py) + LOGGER.info(str(py)) return py From 1f3e9f13654016adabb559018dad45a690bb67ca Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 16:22:34 +0100 Subject: [PATCH 19/25] Cleanup --- utils/benchmarks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 2e407cc93158..db5908473d9e 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -61,6 +61,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}') y.append([name, None, None]) # mAP, t_inference + # Print results LOGGER.info('\n') parse_opt() notebook_init() # print system info From 8310a17a4d077c48e4d55d846e8ab7c1417aec3c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 16:41:22 +0100 Subject: [PATCH 20/25] dir handling fix --- export.py | 6 +++--- utils/benchmarks.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/export.py b/export.py index 37670b428f44..8a483ae878b9 100644 --- a/export.py +++ b/export.py @@ -78,14 +78,14 @@ def export_formats(): x = [['PyTorch', '-', '.pt'], ['TorchScript', 'torchscript', '.torchscript'], ['ONNX', 'onnx', '.onnx'], - ['OpenVINO', 'openvino', '_openvino_model/'], + ['OpenVINO', 'openvino', '_openvino_model'], ['TensorRT', 'engine', '.engine'], ['CoreML', 'coreml', '.mlmodel'], - ['TensorFlow SavedModel', 'saved_model', '_saved_model/'], + ['TensorFlow SavedModel', 'saved_model', '_saved_model'], ['TensorFlow GraphDef', 'pb', '.pb'], ['TensorFlow Lite', 'tflite', '.tflite'], ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite'], - ['TensorFlow.js', 'tfjs', '_web_model/']] + ['TensorFlow.js', 'tfjs', '_web_model']] return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix']) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index db5908473d9e..4da4283e2af2 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -52,7 +52,7 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path for i, (name, f, suffix) in formats.iterrows(): # index, (name, file, suffix) try: w = weights if f == '-' else export.run(weights=weights, imgsz=[imgsz], include=[f], device='cpu')[-1] - assert str(w).endswith(suffix), 'export failed' + assert suffix in str(w), 'export failed' result = val.run(data, w, batch_size, imgsz=imgsz, plots=False, device='cpu', task='benchmark') metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) speeds = result[2] # times (preprocess, inference, postprocess) From b3f423a0cb505a0051e82bbf31374c673760d948 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 11 Feb 2022 15:59:23 +0000 Subject: [PATCH 21/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- utils/benchmarks.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 4da4283e2af2..962df812a9d3 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -37,7 +37,8 @@ sys.path.append(str(ROOT)) # add ROOT to PATH # ROOT = ROOT.relative_to(Path.cwd()) # relative -import export, val +import export +import val from utils import notebook_init from utils.general import LOGGER, print_args From ec6659fb16eb8e89b6640393e01fd486c496203c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 17:05:15 +0100 Subject: [PATCH 22/25] Cleanup --- Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 8a0db815576b..95e2cd4af66d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -40,7 +40,7 @@ ADD https://ultralytics.com/assets/Arial.ttf /root/.config/Ultralytics/ # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t # Kill all -# sudo docker rm $(sudo docker ps -q) +# sudo docker kill $(sudo docker ps -q) # Kill all image-based # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest) From 4733766f1dbf0abca91d389aa930519f28cf3c38 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 17:07:23 +0100 Subject: [PATCH 23/25] Cleanup2 --- utils/general.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/general.py b/utils/general.py index 5f375231e14f..4b7a7c6f0cbf 100755 --- a/utils/general.py +++ b/utils/general.py @@ -422,7 +422,7 @@ def check_dataset(data, autodownload=True): data['names'] = [f'class{i}' for i in range(data['nc'])] # assign class names if missing train, val, test, s = (data.get(x) for x in ('train', 'val', 'test', 'download')) if val: - val = [Path(ROOT / x).resolve() for x in (val if isinstance(val, list) else [val])] # val path + val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path if not all(x.exists() for x in val): LOGGER.info('\nDataset not found, missing paths: %s' % [str(x) for x in val if not x.exists()]) if s and autodownload: # download script From 4efb565098ef7be1fe5ffcf2904ee11f213785f9 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 11 Feb 2022 17:11:33 +0100 Subject: [PATCH 24/25] Cleanup3 --- utils/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/__init__.py b/utils/__init__.py index 2f2eb1dab8e7..4658ed6473cd 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -11,7 +11,7 @@ def notebook_init(verbose=True): import os import shutil - from utils.general import LOGGER, check_requirements, emojis, is_colab + from utils.general import check_requirements, emojis, is_colab from utils.torch_utils import select_device # imports check_requirements(('psutil', 'IPython')) From 6edbf9f0f645cb27bd4d003f34ddeead3c4b36dd Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 12 Feb 2022 15:46:53 +0100 Subject: [PATCH 25/25] Cleanup model_type --- models/common.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/models/common.py b/models/common.py index c17a3d0b7fea..38b94129e274 100644 --- a/models/common.py +++ b/models/common.py @@ -461,12 +461,12 @@ def warmup(self, imgsz=(1, 3, 640, 640), half=False): @staticmethod def model_type(p='path/to/model.pt'): # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx - check_suffix(p, ['.pt', '.torchscript', '.onnx', '.xml', '.engine', '.mlmodel', '.pb', '.tflite', '']) # checks - ends = '.pt', '.torchscript', '.onnx', '.xml', '_openvino_model', '.engine', '.mlmodel', \ - '_saved_model', '.pb', '.tflite', '_edgetpu.tflite', '_web_model' + from export import export_formats + suffixes = list(export_formats().Suffix) + ['.xml'] # export suffixes + check_suffix(p, suffixes) # checks p = Path(p).name # eliminate trailing separators - pt, jit, onnx, xml, xml2, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = (p.endswith(x) for x in ends) - xml |= xml2 # *.xml or *_openvino_model + pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, xml2 = (s in p for s in suffixes) + xml |= xml2 # *_openvino_model or *.xml tflite &= not edgetpu # *.tflite return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs