From fe82e7a0a3b530e843ba9eb706e1517ba1b4515b Mon Sep 17 00:00:00 2001 From: mikel-brostrom Date: Sat, 6 Aug 2022 20:36:40 +0200 Subject: [PATCH 1/4] added export script --- tools/export.py | 224 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 224 insertions(+) create mode 100644 tools/export.py diff --git a/tools/export.py b/tools/export.py new file mode 100644 index 0000000..a613c87 --- /dev/null +++ b/tools/export.py @@ -0,0 +1,224 @@ +import argparse + +import os +# limit the number of cpus used by high performance libraries +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["MKL_NUM_THREADS"] = "1" +os.environ["VECLIB_MAXIMUM_THREADS"] = "1" +os.environ["NUMEXPR_NUM_THREADS"] = "1" + +import sys +import numpy as np +from pathlib import Path +import torch +import pandas as pd +import subprocess +import torch.backends.cudnn as cudnn + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # yolov5 strongsort root directory +WEIGHTS = ROOT / 'weights' + +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +if str(ROOT / 'yolov5') not in sys.path: + sys.path.append(str(ROOT / 'yolov5')) # add yolov5 ROOT to PATH +if str(ROOT / 'strong_sort') not in sys.path: + sys.path.append(str(ROOT / 'strong_sort')) # add strong_sort ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +import logging +from yolov5.models.common import DetectMultiBackend +from yolov5.utils.general import LOGGER, colorstr, check_requirements +from strong_sort.deep.reid.torchreid.utils.feature_extractor import FeatureExtractor +from strong_sort.deep.reid.torchreid.models import build_model +from strong_sort.deep.reid_model_factory import get_model_name + +# remove duplicated stream handler to avoid duplicated logging +logging.getLogger().removeHandler(logging.getLogger().handlers[0]) + +def file_size(path): + # Return file/dir size (MB) + path = Path(path) + if path.is_file(): + return path.stat().st_size / 1E6 + elif path.is_dir(): + return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6 + else: + return 0.0 + + +def export_formats(): + # YOLOv5 export formats + x = [ + ['PyTorch', '-', '.pt', True, True], + ['ONNX', 'onnx', '.onnx', True, True], + ['OpenVINO', 'openvino', '_openvino_model', True, False], + ['TensorFlow Lite', 'tflite', '.tflite', True, False], + ] + return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU']) + + +def export_onnx(model, im, file, opset, train=False, dynamic=True, simplify=False): + # ONNX export + try: + check_requirements(('onnx',)) + import onnx + + f = file.with_suffix('.onnx') + LOGGER.info(f'\nstarting export with onnx {onnx.__version__}...') + + torch.onnx.export( + model.cpu() if dynamic else model, # --dynamic only compatible with cpu + im.cpu() if dynamic else im, + f, + verbose=False, + opset_version=opset, + training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, + do_constant_folding=not train, + input_names=['images'], + output_names=['output'], + dynamic_axes={ + 'images': { + 0: 'batch', + }, # shape(x,3,256,128) + 'output': { + 0: 'batch', + } # shape(x,2048) + } if dynamic else None + ) + # Checks + model_onnx = onnx.load(f) # load onnx model + onnx.checker.check_model(model_onnx) # check onnx model + onnx.save(model_onnx, f) + + # Simplify + if simplify: + try: + cuda = torch.cuda.is_available() + check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1')) + import onnxsim + + LOGGER.info(f'simplifying with onnx-simplifier {onnxsim.__version__}...') + model_onnx, check = onnxsim.simplify( + model_onnx, + dynamic_input_shape=dynamic, + input_shapes={'t0': list(im.shape)} if dynamic else None) + assert check, 'assert check failed' + onnx.save(model_onnx, f) + except Exception as e: + LOGGER.info(f'simplifier failure: {e}') + LOGGER.info(f'export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"run --dynamic ONNX model inference with: 'python detect.py --weights {f}'") + except Exception as e: + LOGGER.info(f'export failure: {e}') + return f + + +def export_openvino(file, dynamic, half, prefix=colorstr('OpenVINO:')): + f = str(file).replace('.onnx', f'_openvino_model{os.sep}') + # YOLOv5 OpenVINO export + try: + check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/ + import openvino.inference_engine as ie + + LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') + f = str(file).replace('.onnx', f'_openvino_model{os.sep}') + dyn_shape = [-1,3,256,128] if dynamic else None + cmd = f"mo \ + --input_model {file} \ + --output_dir {f} \ + --data_type {'FP16' if half else 'FP32'}" + + if dyn_shape is not None: + cmd + f"--input_shape {dyn_shape}" + + subprocess.check_output(cmd.split()) # export + + LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + return f + except Exception as e: + LOGGER.info(f'\n{prefix} export failure: {e}') + return f + + +def export_tflite(file, half, prefix=colorstr('TFLite:')): + # YOLOv5 OpenVINO export + try: + check_requirements(('openvino2tensorflow', 'tensorflow', 'tensorflow_datasets')) # requires openvino-dev: https://pypi.org/project/openvino-dev/ + import openvino.inference_engine as ie + LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') + output = Path(str(file).replace(f'_openvino_model{os.sep}', f'_tflite_model{os.sep}')) + modelxml = list(Path(file).glob('*.xml'))[0] + cmd = f"openvino2tensorflow \ + --model_path {modelxml} \ + --model_output_path {output} \ + --output_pb \ + --output_saved_model \ + --output_no_quant_float32_tflite \ + --output_dynamic_range_quant_tflite" + subprocess.check_output(cmd.split()) # export + + LOGGER.info(f'{prefix} export success, results saved in {output} ({file_size(f):.1f} MB)') + return f + except Exception as e: + LOGGER.info(f'\n{prefix} export failure: {e}') + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="CPHD train") + parser.add_argument( + "-d", + "--dynamic", + action="store_true", + help="dynamic model input", + ) + parser.add_argument( + "-p", + "--weights", + type=Path, + default="./weight/mobilenetv2_x1_0_msmt17.pt", + help="Path to weights", + ) + parser.add_argument( + "-hp", + "--half_precision", + action="store_true", + help="transform model to half precision", + ) + parser.add_argument( + '--imgsz', '--img', '--img-size', + nargs='+', + type=int, + default=[256, 128], + help='image (h, w)' + ) + parser.add_argument('--include', + nargs='+', + default=['onnx', 'openvino', 'tflite'], + help='onnx, openvino, tflite') + args = parser.parse_args() + + # Build model + extractor = FeatureExtractor( + # get rid of dataset information DeepSort model name + model_name=get_model_name(args.weights), + model_path=args.weights, + device=str('cpu') + ) + + include = [x.lower() for x in args.include] # to lowercase + fmts = tuple(export_formats()['Argument'][1:]) # --include arguments + flags = [x in include for x in fmts] + assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}' + onnx, openvino, tflite = flags # export booleans + + im = torch.zeros(1, 3, args.imgsz[0], args.imgsz[1]).to('cpu') # image size(1,3,640,480) BCHW iDetection + if onnx: + f = export_onnx(extractor.model.eval(), im, args.weights, 12, train=False, dynamic=args.dynamic, simplify=True) # opset 12 + if openvino: + f = export_openvino(f, dynamic=args.dynamic, half=False) + if tflite: + export_tflite(f, False) From b49a8a7546e0194e995f369ba30c116ad1cf6f1c Mon Sep 17 00:00:00 2001 From: mikel-brostrom Date: Sat, 6 Aug 2022 21:13:33 +0200 Subject: [PATCH 2/4] cleanup --- requirements.txt | 9 +++++- tools/export.py | 75 +++++++++++++++++++----------------------------- 2 files changed, 37 insertions(+), 47 deletions(-) diff --git a/requirements.txt b/requirements.txt index 6ec2ef0..c61b5f4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,4 +13,11 @@ gdown flake8 yapf isort==4.3.21 -imageio \ No newline at end of file +imageio + +# Export -------------------------------------- +# onnx +# openvino-dev +# openvino2tensorflow +# tensorflow +# tensorflow_datasets \ No newline at end of file diff --git a/tools/export.py b/tools/export.py index a613c87..5e4f72c 100644 --- a/tools/export.py +++ b/tools/export.py @@ -1,42 +1,22 @@ import argparse - import os -# limit the number of cpus used by high performance libraries -os.environ["OMP_NUM_THREADS"] = "1" -os.environ["OPENBLAS_NUM_THREADS"] = "1" -os.environ["MKL_NUM_THREADS"] = "1" -os.environ["VECLIB_MAXIMUM_THREADS"] = "1" -os.environ["NUMEXPR_NUM_THREADS"] = "1" - import sys import numpy as np from pathlib import Path import torch import pandas as pd import subprocess -import torch.backends.cudnn as cudnn -FILE = Path(__file__).resolve() -ROOT = FILE.parents[0] # yolov5 strongsort root directory -WEIGHTS = ROOT / 'weights' -if str(ROOT) not in sys.path: - sys.path.append(str(ROOT)) # add ROOT to PATH -if str(ROOT / 'yolov5') not in sys.path: - sys.path.append(str(ROOT / 'yolov5')) # add yolov5 ROOT to PATH -if str(ROOT / 'strong_sort') not in sys.path: - sys.path.append(str(ROOT / 'strong_sort')) # add strong_sort ROOT to PATH -ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative +import torchreid -import logging -from yolov5.models.common import DetectMultiBackend -from yolov5.utils.general import LOGGER, colorstr, check_requirements -from strong_sort.deep.reid.torchreid.utils.feature_extractor import FeatureExtractor -from strong_sort.deep.reid.torchreid.models import build_model -from strong_sort.deep.reid_model_factory import get_model_name +from torchreid.utils.feature_extractor import FeatureExtractor +from torchreid.models import build_model -# remove duplicated stream handler to avoid duplicated logging -logging.getLogger().removeHandler(logging.getLogger().handlers[0]) +__model_types = [ + 'resnet50', 'mlfn', 'hacnn', 'mobilenetv2_x1_0', 'mobilenetv2_x1_4', + 'osnet_x1_0', 'osnet_x0_75', 'osnet_x0_5', 'osnet_x0_25', + 'osnet_ibn_x1_0', 'osnet_ain_x1_0'] def file_size(path): # Return file/dir size (MB) @@ -49,6 +29,13 @@ def file_size(path): return 0.0 +def get_model_name(model): + model = str(model).rsplit('/', 1)[-1].split('.')[0] + for x in __model_types: + if x in model: + return x + return None + def export_formats(): # YOLOv5 export formats x = [ @@ -63,11 +50,10 @@ def export_formats(): def export_onnx(model, im, file, opset, train=False, dynamic=True, simplify=False): # ONNX export try: - check_requirements(('onnx',)) import onnx f = file.with_suffix('.onnx') - LOGGER.info(f'\nstarting export with onnx {onnx.__version__}...') + print(f'\nStarting export with onnx {onnx.__version__}...') torch.onnx.export( model.cpu() if dynamic else model, # --dynamic only compatible with cpu @@ -97,10 +83,9 @@ def export_onnx(model, im, file, opset, train=False, dynamic=True, simplify=Fals if simplify: try: cuda = torch.cuda.is_available() - check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1')) import onnxsim - LOGGER.info(f'simplifying with onnx-simplifier {onnxsim.__version__}...') + print(f'simplifying with onnx-simplifier {onnxsim.__version__}...') model_onnx, check = onnxsim.simplify( model_onnx, dynamic_input_shape=dynamic, @@ -108,22 +93,21 @@ def export_onnx(model, im, file, opset, train=False, dynamic=True, simplify=Fals assert check, 'assert check failed' onnx.save(model_onnx, f) except Exception as e: - LOGGER.info(f'simplifier failure: {e}') - LOGGER.info(f'export success, saved as {f} ({file_size(f):.1f} MB)') - LOGGER.info(f"run --dynamic ONNX model inference with: 'python detect.py --weights {f}'") + print(f'simplifier failure: {e}') + print(f'export success, saved as {f} ({file_size(f):.1f} MB)') + print(f"run --dynamic ONNX model inference with: 'python detect.py --weights {f}'") except Exception as e: - LOGGER.info(f'export failure: {e}') + print(f'export failure: {e}') return f -def export_openvino(file, dynamic, half, prefix=colorstr('OpenVINO:')): +def export_openvino(file, dynamic, half): f = str(file).replace('.onnx', f'_openvino_model{os.sep}') # YOLOv5 OpenVINO export try: - check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/ import openvino.inference_engine as ie - LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') + print(f'\nStarting export with openvino {ie.__version__}...') f = str(file).replace('.onnx', f'_openvino_model{os.sep}') dyn_shape = [-1,3,256,128] if dynamic else None cmd = f"mo \ @@ -136,19 +120,18 @@ def export_openvino(file, dynamic, half, prefix=colorstr('OpenVINO:')): subprocess.check_output(cmd.split()) # export - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + print(f'Export success, saved as {f} ({file_size(f):.1f} MB)') return f except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + print(f'\nExport failure: {e}') return f -def export_tflite(file, half, prefix=colorstr('TFLite:')): +def export_tflite(file, half): # YOLOv5 OpenVINO export try: - check_requirements(('openvino2tensorflow', 'tensorflow', 'tensorflow_datasets')) # requires openvino-dev: https://pypi.org/project/openvino-dev/ import openvino.inference_engine as ie - LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') + print(f'\nStarting export with openvino {ie.__version__}...') output = Path(str(file).replace(f'_openvino_model{os.sep}', f'_tflite_model{os.sep}')) modelxml = list(Path(file).glob('*.xml'))[0] cmd = f"openvino2tensorflow \ @@ -160,10 +143,10 @@ def export_tflite(file, half, prefix=colorstr('TFLite:')): --output_dynamic_range_quant_tflite" subprocess.check_output(cmd.split()) # export - LOGGER.info(f'{prefix} export success, results saved in {output} ({file_size(f):.1f} MB)') + print(f'Export success, results saved in {output} ({file_size(f):.1f} MB)') return f except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + print(f'\nExport failure: {e}') if __name__ == "__main__": @@ -179,7 +162,7 @@ if __name__ == "__main__": "-p", "--weights", type=Path, - default="./weight/mobilenetv2_x1_0_msmt17.pt", + default="./mobilenetv2_x1_0_msmt17.pt", help="Path to weights", ) parser.add_argument( From ac8008688206b0203b2e7d1fcb6d2e6e6b2267ce Mon Sep 17 00:00:00 2001 From: mikel-brostrom Date: Sat, 6 Aug 2022 21:14:46 +0200 Subject: [PATCH 3/4] delete unused imports --- tools/export.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tools/export.py b/tools/export.py index 5e4f72c..28dfde2 100644 --- a/tools/export.py +++ b/tools/export.py @@ -7,9 +7,6 @@ import torch import pandas as pd import subprocess - -import torchreid - from torchreid.utils.feature_extractor import FeatureExtractor from torchreid.models import build_model From 9e4071b5c8da9744fa0f3f712257c70e9ade24e1 Mon Sep 17 00:00:00 2001 From: Mike Date: Sat, 6 Aug 2022 21:51:00 +0200 Subject: [PATCH 4/4] Update requirements.txt --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index c61b5f4..71da6fe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,7 +17,8 @@ imageio # Export -------------------------------------- # onnx +# onnx-simplified # openvino-dev # openvino2tensorflow # tensorflow -# tensorflow_datasets \ No newline at end of file +# tensorflow_datasets