diff --git a/README.md b/README.md index 457e2a5..9635023 100644 --- a/README.md +++ b/README.md @@ -2,16 +2,14 @@ NVIDIA DeepStream SDK 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 configuration for YOLO models -------------------------------------- -### **Big update on DeepStream-Yolo** -------------------------------------- +-------------------------------------------------------------------------------------------------- ### Important: please generate the ONNX model and the TensorRT engine again with the updated files -------------------------------------- +-------------------------------------------------------------------------------------------------- ### Future updates * DeepStream tutorials -* Dynamic batch-size + * Updated INT8 calibration * Support for segmentation models * Support for classification models @@ -24,6 +22,7 @@ NVIDIA DeepStream SDK 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 configuration for YOLO mod * **Support for Darknet YOLO models (YOLOv4, etc) using cfg and weights conversion with GPU post-processing** * **Support for YOLO-NAS, PPYOLOE+, PPYOLOE, DAMO-YOLO, YOLOX, YOLOR, YOLOv8, YOLOv7, YOLOv6 and YOLOv5 using ONNX conversion with GPU post-processing** * **Add GPU bbox parser (it is slightly slower than CPU bbox parser on V100 GPU tests)** +* **Dynamic batch-size for ONNX exported models (YOLO-NAS, PPYOLOE+, PPYOLOE, DAMO-YOLO, YOLOX, YOLOR, YOLOv8, YOLOv7, YOLOv6 and YOLOv5)** ## diff --git a/docs/DAMOYOLO.md b/docs/DAMOYOLO.md index a8be40c..2c352ea 100644 --- a/docs/DAMOYOLO.md +++ b/docs/DAMOYOLO.md @@ -40,7 +40,7 @@ wget https://idstcv.oss-cn-zhangjiakou.aliyuncs.com/DAMO-YOLO/release_model/clea Generate the ONNX model file (example for DAMO-YOLO-S*) ``` -python3 export_damoyolo.py -w damoyolo_tinynasL25_S_477.pth -c configs/damoyolo_tinynasL25_S.py --simplify +python3 export_damoyolo.py -w damoyolo_tinynasL25_S_477.pth -c configs/damoyolo_tinynasL25_S.py --simplify --dynamic ``` **NOTE**: To change the inference size (defaut: 640) @@ -66,7 +66,7 @@ or #### 5. Copy generated files -Copy the generated ONNX model file to the `DeepStream-Yolo` folder. +Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder. ## diff --git a/docs/PPYOLOE.md b/docs/PPYOLOE.md index 1fd6065..bb4672b 100644 --- a/docs/PPYOLOE.md +++ b/docs/PPYOLOE.md @@ -43,7 +43,7 @@ python3 export_ppyoloe.py -w ppyoloe_plus_crn_s_80e_coco.pdparams -c configs/ppy #### 5. Copy generated files -Copy the generated ONNX model file to the `DeepStream-Yolo` folder. +Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder. ## diff --git a/docs/YOLONAS.md b/docs/YOLONAS.md index eb94fa6..3b7badf 100644 --- a/docs/YOLONAS.md +++ b/docs/YOLONAS.md @@ -43,7 +43,7 @@ wget https://sghub.deci.ai/models/yolo_nas_s_coco.pth Generate the ONNX model file (example for YOLO-NAS S) ``` -python3 export_yolonas.py -m yolo_nas_s -w yolo_nas_s_coco.pth --simplify +python3 export_yolonas.py -m yolo_nas_s -w yolo_nas_s_coco.pth --simplify --dynamic ``` **NOTE**: Model names @@ -85,7 +85,7 @@ or -s 1280 1280 ``` -#### 5. Copy generated files +#### 5. Copy generated file Copy the generated ONNX model file to the `DeepStream-Yolo` folder. diff --git a/docs/YOLOR.md b/docs/YOLOR.md index 0712c7f..600c42f 100644 --- a/docs/YOLOR.md +++ b/docs/YOLOR.md @@ -44,7 +44,7 @@ Generate the ONNX model file Example for YOLOR-CSP ``` - python3 export_yolor.py -w yolor_csp.pt -c cfg/yolor_csp.cfg --simplify + python3 export_yolor.py -w yolor_csp.pt -c cfg/yolor_csp.cfg --simplify --dynamic ``` - Paper branch @@ -52,7 +52,7 @@ Generate the ONNX model file Example for YOLOR-P6 ``` - python3 export_yolor.py -w yolor-p6.pt --simplify + python3 export_yolor.py -w yolor-p6.pt --simplify --dynamic ``` **NOTE**: To convert a P6 model @@ -84,7 +84,7 @@ or #### 5. Copy generated files -Copy the generated ONNX model file to the `DeepStream-Yolo` folder +Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder ## diff --git a/docs/YOLOX.md b/docs/YOLOX.md index da8835e..ffabc4d 100644 --- a/docs/YOLOX.md +++ b/docs/YOLOX.md @@ -43,10 +43,10 @@ wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yo Generate the ONNX model file (example for YOLOX-s) ``` -python3 export_yolox.py -w yolox_s.pth -c exps/default/yolox_s.py --simplify +python3 export_yolox.py -w yolox_s.pth -c exps/default/yolox_s.py --simplify --dynamic ``` -#### 5. Copy generated files +#### 5. Copy generated file Copy the generated ONNX model file to the `DeepStream-Yolo` folder. diff --git a/docs/YOLOv5.md b/docs/YOLOv5.md index 86de7a5..6e06eeb 100644 --- a/docs/YOLOv5.md +++ b/docs/YOLOv5.md @@ -44,7 +44,7 @@ wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt Generate the ONNX model file (example for YOLOv5s) ``` -python3 export_yoloV5.py -w yolov5s.pt --simplify +python3 export_yoloV5.py -w yolov5s.pt --simplify --dynamic ``` **NOTE**: To convert a P6 model @@ -76,7 +76,7 @@ or #### 5. Copy generated files -Copy the generated ONNX model file to the `DeepStream-Yolo` folder. +Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder. ## diff --git a/docs/YOLOv6.md b/docs/YOLOv6.md index bc5b517..7d77726 100644 --- a/docs/YOLOv6.md +++ b/docs/YOLOv6.md @@ -44,7 +44,7 @@ wget https://github.com/meituan/YOLOv6/releases/download/0.4.0/yolov6s.pt Generate the ONNX model file (example for YOLOv6-S 4.0) ``` -python3 export_yoloV6.py -w yolov6s.pt --simplify +python3 export_yoloV6.py -w yolov6s.pt --simplify --dynamic ``` **NOTE**: To convert a P6 model @@ -74,7 +74,7 @@ or -s 1280 1280 ``` -#### 5. Copy generated files +#### 5. Copy generated file Copy the generated ONNX model file to the `DeepStream-Yolo` folder. diff --git a/docs/YOLOv7.md b/docs/YOLOv7.md index d448e51..b046520 100644 --- a/docs/YOLOv7.md +++ b/docs/YOLOv7.md @@ -46,7 +46,7 @@ wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt Generate the ONNX model file (example for YOLOv7) ``` -python3 export_yoloV7.py -w yolov7.pt --simplify +python3 export_yoloV7.py -w yolov7.pt --simplify --dynamic ``` **NOTE**: To convert a P6 model @@ -78,7 +78,7 @@ or #### 6. Copy generated files -Copy the generated ONNX model file to the `DeepStream-Yolo` folder. +Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder. ## diff --git a/docs/YOLOv8.md b/docs/YOLOv8.md index 26cb69c..0239573 100644 --- a/docs/YOLOv8.md +++ b/docs/YOLOv8.md @@ -43,7 +43,7 @@ wget https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt Generate the ONNX model file (example for YOLOv8s) ``` -python3 export_yoloV8.py -w yolov8s.pt --simplify +python3 export_yoloV8.py -w yolov8s.pt --simplify --dynamic ``` **NOTE**: To change the inference size (defaut: 640) @@ -69,7 +69,7 @@ or #### 5. Copy generated files -Copy the generated ONNX model file to the `DeepStream-Yolo` folder. +Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder. ## diff --git a/utils/export_damoyolo.py b/utils/export_damoyolo.py index 68a6744..ff0ebbf 100644 --- a/utils/export_damoyolo.py +++ b/utils/export_damoyolo.py @@ -46,9 +46,21 @@ def damoyolo_export(weights, config_file, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening DAMO-YOLO model') + device = torch.device('cpu') cfg, model = damoyolo_export(args.weights, args.config, device) + if len(cfg.dataset['class_names']) > 0: + print('Creating labels.txt file') + f = open('labels.txt', 'w') + for name in cfg.dataset['class_names']: + f.write(name + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = args.size * 2 if len(args.size) == 1 else args.size @@ -56,15 +68,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = cfg.miscs['exp_name'] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('Exporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream DAMO-YOLO conversion') @@ -73,6 +99,7 @@ def parse_args(): parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])') parser.add_argument('--opset', type=int, default=11, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_ppyoloe.py b/utils/export_ppyoloe.py index 1c2fca6..a03c26d 100644 --- a/utils/export_ppyoloe.py +++ b/utils/export_ppyoloe.py @@ -8,6 +8,7 @@ from ppdet.utils.check import check_version, check_config from ppdet.utils.cli import ArgsParser from ppdet.engine import Trainer from ppdet.slim import build_slim_model +from ppdet.data.source.category import get_categories class DeepStreamOutput(nn.Layer): @@ -39,13 +40,26 @@ def ppyoloe_export(FLAGS): os.makedirs('.tmp') static_model, _ = trainer._get_infer_cfg_and_input_spec('.tmp') os.system('rm -r .tmp') - return cfg, static_model + return trainer.cfg, static_model def main(FLAGS): + print('\nStarting: %s' % FLAGS.weights) + + print('\nOpening PPYOLOE model\n') + paddle.set_device('cpu') cfg, model = ppyoloe_export(FLAGS) + anno_file = cfg['TestDataset'].get_anno() + if os.path.isfile(anno_file): + _, catid2name = get_categories(cfg['metric'], anno_file, 'detection_arch') + print('\nCreating labels.txt file') + f = open('labels.txt', 'w') + for name in catid2name.values(): + f.write(str(name) + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = [cfg.eval_height, cfg.eval_width] @@ -55,14 +69,18 @@ def main(FLAGS): onnx_input_im['scale_factor'] = paddle.static.InputSpec(shape=[None, 2], dtype='float32', name='scale_factor') onnx_output_file = cfg.filename + '.onnx' + print('\nExporting the model to ONNX\n') paddle.onnx.export(model, cfg.filename, input_spec=[onnx_input_im], opset_version=FLAGS.opset) if FLAGS.simplify: + print('\nSimplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('\nDone: %s\n' % onnx_output_file) + def parse_args(): parser = ArgsParser() diff --git a/utils/export_yoloV5.py b/utils/export_yoloV5.py index 4edbe78..ca2703c 100644 --- a/utils/export_yoloV5.py +++ b/utils/export_yoloV5.py @@ -41,9 +41,21 @@ def yolov5_export(weights, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOv5 model\n') + device = select_device('cpu') model = yolov5_export(args.weights, device) + if len(model.names.keys()) > 0: + print('\nCreating labels.txt file') + f = open('labels.txt', 'w') + for name in model.names.values(): + f.write(name + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = args.size * 2 if len(args.size) == 1 else args.size @@ -54,15 +66,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOv5 conversion') @@ -71,6 +97,7 @@ def parse_args(): parser.add_argument('--p6', action='store_true', help='P6 model') parser.add_argument('--opset', type=int, default=17, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_yoloV6.py b/utils/export_yoloV6.py index dc51a23..6b224f3 100644 --- a/utils/export_yoloV6.py +++ b/utils/export_yoloV6.py @@ -51,6 +51,11 @@ def yolov6_export(weights, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOv6 model\n') + device = torch.device('cpu') model = yolov6_export(args.weights, device) @@ -64,15 +69,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOv6 conversion') @@ -81,6 +100,7 @@ def parse_args(): parser.add_argument('--p6', action='store_true', help='P6 model') parser.add_argument('--opset', type=int, default=13, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_yoloV7.py b/utils/export_yoloV7.py index 06e452f..d33b001 100644 --- a/utils/export_yoloV7.py +++ b/utils/export_yoloV7.py @@ -45,9 +45,21 @@ def yolov7_export(weights, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOv7 model\n') + device = select_device('cpu') model = yolov7_export(args.weights, device) + if len(model.names) > 0: + print('\nCreating labels.txt file') + f = open('labels.txt', 'w') + for name in model.names: + f.write(name + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = args.size * 2 if len(args.size) == 1 else args.size @@ -58,15 +70,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOv7 conversion') @@ -75,6 +101,7 @@ def parse_args(): parser.add_argument('--p6', action='store_true', help='P6 model') parser.add_argument('--opset', type=int, default=12, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_yoloV7_u6.py b/utils/export_yoloV7_u6.py index bd7f399..d1fd5f5 100644 --- a/utils/export_yoloV7_u6.py +++ b/utils/export_yoloV7_u6.py @@ -40,9 +40,21 @@ def yolov7_u6_export(weights, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOv7_u6 model\n') + device = select_device('cpu') model = yolov7_u6_export(args.weights, device) + if len(model.names.keys()) > 0: + print('\nCreating labels.txt file') + f = open('labels.txt', 'w') + for name in model.names.values(): + f.write(name + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = args.size * 2 if len(args.size) == 1 else args.size @@ -50,15 +62,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOv7-u6 conversion') @@ -66,6 +92,7 @@ def parse_args(): parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])') parser.add_argument('--opset', type=int, default=12, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_yoloV8.py b/utils/export_yoloV8.py index 907f36c..12d94f1 100644 --- a/utils/export_yoloV8.py +++ b/utils/export_yoloV8.py @@ -48,9 +48,21 @@ def yolov8_export(weights, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOv8 model\n') + device = select_device('cpu') model = yolov8_export(args.weights, device) + if len(model.names.keys()) > 0: + print('\nCreating labels.txt file') + f = open('labels.txt', 'w') + for name in model.names.values(): + f.write(name + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = args.size * 2 if len(args.size) == 1 else args.size @@ -58,15 +70,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOv8 conversion') @@ -74,6 +100,7 @@ def parse_args(): parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])') parser.add_argument('--opset', type=int, default=16, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_yolonas.py b/utils/export_yolonas.py index 3cf37b0..d5ca4b8 100644 --- a/utils/export_yolonas.py +++ b/utils/export_yolonas.py @@ -34,6 +34,11 @@ def yolonas_export(model_name, weights, num_classes, size): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLO-NAS model\n') + device = torch.device('cpu') model = yolonas_export(args.model, args.weights, args.classes, args.size) @@ -44,15 +49,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLO-NAS conversion') @@ -62,6 +81,7 @@ def parse_args(): parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])') parser.add_argument('--opset', type=int, default=14, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if args.model == '': raise SystemExit('Invalid model name') diff --git a/utils/export_yolor.py b/utils/export_yolor.py index b20c1e2..33d6960 100644 --- a/utils/export_yolor.py +++ b/utils/export_yolor.py @@ -57,9 +57,21 @@ def yolor_export(weights, cfg, size, device): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOR model\n') + device = torch.device('cpu') model = yolor_export(args.weights, args.cfg, args.size, device) + if hasattr(model, 'names') and len(model.names) > 0: + print('\nCreating labels.txt file') + f = open('labels.txt', 'w') + for name in model.names: + f.write(name + '\n') + f.close() + model = nn.Sequential(model, DeepStreamOutput()) img_size = args.size * 2 if len(args.size) == 1 else args.size @@ -70,15 +82,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('\nExporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOR conversion') @@ -88,6 +114,7 @@ def parse_args(): parser.add_argument('--p6', action='store_true', help='P6 model') parser.add_argument('--opset', type=int, default=12, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file') diff --git a/utils/export_yolox.py b/utils/export_yolox.py index f51c61c..3cdd8f3 100644 --- a/utils/export_yolox.py +++ b/utils/export_yolox.py @@ -42,6 +42,11 @@ def yolox_export(weights, exp_file): def main(args): suppress_warnings() + + print('\nStarting: %s' % args.weights) + + print('Opening YOLOX model') + device = torch.device('cpu') model, exp = yolox_export(args.weights, args.exp) @@ -52,15 +57,29 @@ def main(args): onnx_input_im = torch.zeros(1, 3, *img_size).to(device) onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx' + dynamic_axes = { + 'input': { + 0: 'batch' + }, + 'output': { + 0: 'batch' + } + } + + print('Exporting the model to ONNX') torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, - do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes=None) + do_constant_folding=True, input_names=['input'], output_names=['output'], + dynamic_axes=dynamic_axes if args.dynamic else None) if args.simplify: + print('Simplifying the ONNX model') import onnxsim model_onnx = onnx.load(onnx_output_file) model_onnx, _ = onnxsim.simplify(model_onnx) onnx.save(model_onnx, onnx_output_file) + print('Done: %s\n' % onnx_output_file) + def parse_args(): parser = argparse.ArgumentParser(description='DeepStream YOLOX conversion') @@ -68,6 +87,7 @@ def parse_args(): parser.add_argument('-c', '--exp', required=True, help='Input exp (.py) file path (required)') parser.add_argument('--opset', type=int, default=11, help='ONNX opset version') parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') + parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') args = parser.parse_args() if not os.path.isfile(args.weights): raise SystemExit('Invalid weights file')