From 698a5d7f26002e7b0b0d535d981c2b92f25bc76e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 9 Apr 2022 01:32:16 +0200 Subject: [PATCH] Add `python benchmarks.py --test` for export-only (#7350) * Test exports * Fix precommit --- utils/benchmarks.py | 44 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/utils/benchmarks.py b/utils/benchmarks.py index 36e827848584..1c1bb7a8f2ed 100644 --- a/utils/benchmarks.py +++ b/utils/benchmarks.py @@ -52,20 +52,26 @@ def run( data=ROOT / 'data/coco128.yaml', # dataset.yaml path device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu half=False, # use FP16 half-precision inference + test=False, # test exports only ): y, t = [], time.time() formats = export.export_formats() device = select_device(device) for i, (name, f, suffix, gpu) in formats.iterrows(): # index, (name, file, suffix, gpu-capable) try: - assert i < 9, 'Edge TPU and TF.js not supported' + assert i != 9, 'Edge TPU not supported' + assert i != 10, 'TF.js not supported' if device.type != 'cpu': assert gpu, f'{name} inference not supported on GPU' + + # Export if f == '-': w = weights # PyTorch format else: w = export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # all others assert suffix in str(w), 'export failed' + + # Validate result = val.run(data, w, batch_size, imgsz, plots=False, device=device, task='benchmark', half=half) metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls)) speeds = result[2] # times (preprocess, inference, postprocess) @@ -78,8 +84,39 @@ def run( LOGGER.info('\n') parse_opt() notebook_init() # print system info - py = pd.DataFrame(y, columns=['Format', 'mAP@0.5:0.95', 'Inference time (ms)']) + py = pd.DataFrame(y, columns=['Format', 'mAP@0.5:0.95', 'Inference time (ms)'] if map else ['Format', 'Export', '']) LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)') + LOGGER.info(str(py if map else py.iloc[:, :2])) + return py + + +def test( + weights=ROOT / 'yolov5s.pt', # weights path + imgsz=640, # inference size (pixels) + batch_size=1, # batch size + data=ROOT / 'data/coco128.yaml', # dataset.yaml path + device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu + half=False, # use FP16 half-precision inference + test=False, # test exports only +): + y, t = [], time.time() + formats = export.export_formats() + device = select_device(device) + for i, (name, f, suffix, gpu) in formats.iterrows(): # index, (name, file, suffix, gpu-capable) + try: + w = weights if f == '-' else \ + export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # weights + assert suffix in str(w), 'export failed' + y.append([name, True]) + except Exception: + y.append([name, False]) # mAP, t_inference + + # Print results + LOGGER.info('\n') + parse_opt() + notebook_init() # print system info + py = pd.DataFrame(y, columns=['Format', 'Export']) + LOGGER.info(f'\nExports complete ({time.time() - t:.2f}s)') LOGGER.info(str(py)) return py @@ -92,13 +129,14 @@ def parse_opt(): parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') + parser.add_argument('--test', action='store_true', help='test exports only') opt = parser.parse_args() print_args(vars(opt)) return opt def main(opt): - run(**vars(opt)) + test(**vars(opt)) if opt.test else run(**vars(opt)) if __name__ == "__main__":