├── .github └── ISSUE_TEMPLATE │ ├── config.yml │ ├── feature_request.yml │ └── question.yml ├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── README_cn.md ├── assets ├── banner-YOLO.png ├── image3.jpg ├── picture.png ├── speed_comparision_v2.png ├── speed_comparision_v3.png ├── train_batch.jpg ├── voc_loss_curve.jpg ├── wechat_qrcode.png ├── yolov5s.jpg ├── yolov6s.jpg └── yoloxs.jpg ├── configs ├── base │ ├── README.md │ ├── README_cn.md │ ├── yolov6l_base.py │ ├── yolov6l_base_finetune.py │ ├── yolov6m_base.py │ ├── yolov6m_base_finetune.py │ ├── yolov6n_base.py │ ├── yolov6n_base_finetune.py │ ├── yolov6s_base.py │ └── yolov6s_base_finetune.py ├── experiment │ ├── eval_640_repro.py │ ├── yolov6n_with_eval_params.py │ ├── yolov6s_csp_scaled.py │ ├── yolov6t.py │ ├── yolov6t_csp_scaled.py │ └── yolov6t_finetune.py ├── qarepvgg │ ├── README.md │ ├── yolov6m_qa.py │ ├── yolov6n_qa.py │ └── yolov6s_qa.py ├── repopt │ ├── yolov6_tiny_hs.py │ ├── yolov6_tiny_opt.py │ ├── yolov6_tiny_opt_qat.py │ ├── yolov6n_hs.py │ ├── yolov6n_opt.py │ ├── yolov6n_opt_qat.py │ ├── yolov6s_hs.py │ ├── yolov6s_opt.py │ └── yolov6s_opt_qat.py ├── yolov6l.py ├── yolov6l6.py ├── yolov6l6_finetune.py ├── yolov6l_finetune.py ├── yolov6m.py ├── yolov6m6.py ├── yolov6m6_finetune.py ├── yolov6m_finetune.py ├── yolov6n.py ├── yolov6n6.py ├── yolov6n6_finetune.py ├── yolov6n_finetune.py ├── yolov6s.py ├── yolov6s6.py ├── yolov6s6_finetune.py └── yolov6s_finetune.py ├── data ├── coco.yaml ├── dataset.yaml ├── images │ ├── image1.jpg │ ├── image2.jpg │ └── image3.jpg └── voc.yaml ├── deploy ├── ONNX │ ├── OpenCV │ │ ├── README.md │ │ ├── coco.names │ │ ├── sample.jpg │ │ ├── yolo.py │ │ ├── yolo_video.py │ │ ├── yolov5 │ │ │ ├── CMakeLists.txt │ │ │ └── yolov5.cpp │ │ ├── yolov6 │ │ │ ├── CMakeLists.txt │ │ │ └── yolov6.cpp │ │ ├── yolox.py │ │ └── yolox │ │ │ ├── CMakeLists.txt │ │ │ └── yolox.cpp │ ├── README.md │ ├── YOLOv6-Dynamic-Batch-onnxruntime.ipynb │ ├── YOLOv6-Dynamic-Batch-tensorrt.ipynb │ ├── eval_trt.py │ └── export_onnx.py ├── OpenVINO │ ├── README.md │ └── export_openvino.py ├── RKNN │ ├── RKOPT_README.md │ ├── RKOPT_README_cn.md │ └── export_onnx_for_rknn.py └── TensorRT │ ├── CMakeLists.txt │ ├── README.md │ ├── calibrator.py │ ├── eval_yolo_trt.py │ ├── logging.h │ ├── onnx_to_trt.py │ ├── tensorrt_processor.py │ ├── visualize.py │ └── yolov6.cpp ├── docs ├── Test_speed.md ├── Train_coco_data.md ├── Train_custom_data.md ├── Tutorial of Quantization.md ├── tutorial_repopt.md └── tutorial_voc.ipynb ├── hubconf.py ├── inference.ipynb ├── requirements.txt ├── tools ├── eval.py ├── infer.py ├── partial_quantization │ ├── README.md │ ├── eval.py │ ├── eval.yaml │ ├── partial_quant.py │ ├── ptq.py │ ├── sensitivity_analyse.py │ └── utils.py ├── qat │ ├── README.md │ ├── onnx_utils.py │ ├── qat_export.py │ └── qat_utils.py ├── quantization │ ├── mnn │ │ └── README.md │ ├── ppq │ │ ├── ProgramEntrance.py │ │ └── write_qparams_onnx2trt.py │ └── tensorrt │ │ ├── post_training │ │ ├── Calibrator.py │ │ ├── LICENSE │ │ ├── README.md │ │ ├── onnx_to_tensorrt.py │ │ └── quant.sh │ │ ├── requirements.txt │ │ └── training_aware │ │ └── QAT_quantizer.py └── train.py ├── turtorial.ipynb └── yolov6 ├── __init__.py ├── assigners ├── __init__.py ├── anchor_generator.py ├── assigner_utils.py ├── atss_assigner.py ├── iou2d_calculator.py └── tal_assigner.py ├── core ├── engine.py ├── evaler.py └── inferer.py ├── data ├── data_augment.py ├── data_load.py ├── datasets.py ├── vis_dataset.py └── voc2yolo.py ├── layers ├── common.py └── dbb_transforms.py ├── models ├── efficientrep.py ├── effidehead.py ├── end2end.py ├── heads │ ├── effidehead_distill_ns.py │ └── effidehead_fuseab.py ├── losses │ ├── loss.py │ ├── loss_distill.py │ ├── loss_distill_ns.py │ └── loss_fuseab.py ├── reppan.py └── yolo.py ├── solver └── build.py └── utils ├── Arial.ttf ├── RepOptimizer.py ├── checkpoint.py ├── config.py ├── ema.py ├── envs.py ├── events.py ├── figure_iou.py ├── general.py ├── metrics.py ├── nms.py └── torch_utils.py /.github/ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- 1 | blank_issues_enabled: false 2 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.yml: -------------------------------------------------------------------------------- 1 | name: 🚀 Feature Request 2 | description: Suggest a YOLOv6 idea 3 | # title: " " 4 | labels: [enhancement] 5 | body: 6 | - type: markdown 7 | attributes: 8 | value: | 9 | Thank you for submitting a YOLOv6 Feature Request! 10 | 11 | - type: checkboxes 12 | attributes: 13 | label: Search before asking 14 | description: > 15 | Please search the [issues](https://github.com/meituan/YOLOv6/issues) to see if a similar feature request already exists. 16 | options: 17 | - label: > 18 | I have searched the YOLOv6 [issues](https://github.com/meituan/YOLOv6/issues) and found no similar feature requests. 19 | required: true 20 | 21 | - type: textarea 22 | attributes: 23 | label: Description 24 | description: A short description of your feature. 25 | placeholder: | 26 | What new feature would you like to see in YOLOv6? （你希望YOLOv6可以支持哪些新功能？） 27 | validations: 28 | required: true 29 | 30 | - type: textarea 31 | attributes: 32 | label: Use case 33 | description: | 34 | Describe the use case of your feature request. It will help us understand and prioritize the feature request. 35 | placeholder: | 36 | How would this feature be used, and who would use it?（请描述一下这个新功能的应用场景？） 37 | 38 | - type: textarea 39 | attributes: 40 | label: Additional 41 | description: Anything else you would like to share? 42 | 43 | - type: checkboxes 44 | attributes: 45 | label: Are you willing to submit a PR? 46 | description: > 47 | (Optional) We encourage you to submit a [Pull Request](https://github.com/meituan/YOLOv6/pulls) (PR) to help improve YOLOv6 for everyone, especially if you have a good understanding of how to implement a fix or feature. 48 | options: 49 | - label: Yes I'd like to help by submitting a PR! 50 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.yml: -------------------------------------------------------------------------------- 1 | name: ❓ Question 2 | description: Ask a YOLOv6 question 3 | # title: " " 4 | labels: [question] 5 | body: 6 | - type: markdown 7 | attributes: 8 | value: | 9 | Thanks for your attention. We will try our best to solve your problem, but more concrete information is necessary to reproduce your problem. 10 | - type: checkboxes 11 | attributes: 12 | label: Before Asking 13 | description: > 14 | Please check and try following methods to solve it by yourself 15 | options: 16 | - label: > 17 | I have read the [README](https://github.com/meituan/YOLOv6/blob/main/README.md) carefully. 18 | 我已经仔细阅读了README上的操作指引。 19 | required: true 20 | - label: > 21 | I want to train my custom dataset, and I have read the [tutorials for training your custom data](https://github.com/meituan/YOLOv6/blob/main/docs/Train_custom_data.md) carefully and organize my dataset correctly; 22 | (FYI: We recommand you to apply the config files of xx_finetune.py.) 23 | 我想训练自定义数据集，我已经仔细阅读了训练自定义数据的教程，以及按照正确的目录结构存放数据集。（FYI: 我们推荐使用xx_finetune.py等配置文件训练自定义数据集。） 24 | required: False 25 | - label: > 26 | I have pulled the latest code of main branch to run again and the problem still existed. 27 | 我已经拉取了主分支上最新的代码，重新运行之后，问题仍不能解决。 28 | required: true 29 | 30 | 31 | - type: checkboxes 32 | attributes: 33 | label: Search before asking 34 | description: > 35 | Please search the [issues](https://github.com/meituan/YOLOv6/issues) to see if a similar question already exists. 36 | options: 37 | - label: > 38 | I have searched the YOLOv6 [issues](https://github.com/meituan/YOLOv6/issues) and found no similar questions. 39 | required: true 40 | 41 | - type: textarea 42 | attributes: 43 | label: Question 44 | description: What is your question? 45 | placeholder: | 46 | 💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks, training commands etc.) to receive the most helpful response. 47 | (请仔细阅读上面的信息先进行问题排查，如果仍不能解决您的问题，请将问题尽可能地描述详细，以及提供相关命令、超参配置、报错日志等信息或截图，以便更快地定位和解决问题。) 48 | validations: 49 | required: true 50 | 51 | - type: textarea 52 | attributes: 53 | label: Additional 54 | description: Anything else you would like to share? 55 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | **/*.pyc 6 | 7 | # C extensions 8 | 9 | # Distribution / packaging 10 | 11 | .Python 12 | videos/ 13 | build/ 14 | runs/ 15 | weights/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | .hypothesis/ 50 | .pytest_cache/ 51 | 52 | # Translations 53 | *.mo 54 | *.pot 55 | 56 | # Django stuff: 57 | *.log 58 | local_settings.py 59 | db.sqlite3 60 | 61 | # Flask stuff: 62 | instance/ 63 | .webassets-cache 64 | 65 | # Scrapy stuff: 66 | .scrapy 67 | 68 | # Sphinx documentation 69 | docs/_build/ 70 | 71 | # PyBuilder 72 | target/ 73 | 74 | # Jupyter Notebook 75 | .ipynb_checkpoints 76 | 77 | # pyenv 78 | .python-version 79 | 80 | # celery beat schedule file 81 | celerybeat-schedule 82 | 83 | # SageMath parsed files 84 | *.sage.py 85 | 86 | # Environments 87 | .env 88 | .venv 89 | env/ 90 | venv/ 91 | ENV/ 92 | env.bak/ 93 | venv.bak/ 94 | 95 | # Spyder project settings 96 | .spyderproject 97 | .spyproject 98 | 99 | # Rope project settings 100 | .ropeproject 101 | 102 | # custom 103 | .DS_Store 104 | 105 | # Pytorch 106 | *.pth 107 | 108 | #vscode 109 | .vscode/* 110 | 111 | #user scripts 112 | *.sh 113 | 114 | # model files 115 | *.onnx 116 | *.pt 117 | *.engine 118 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v4.3.0 4 | hooks: 5 | - id: check-yaml 6 | - id: end-of-file-fixer 7 | - id: trailing-whitespace 8 | -------------------------------------------------------------------------------- /assets/banner-YOLO.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/banner-YOLO.png -------------------------------------------------------------------------------- /assets/image3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/image3.jpg -------------------------------------------------------------------------------- /assets/picture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/picture.png -------------------------------------------------------------------------------- /assets/speed_comparision_v2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/speed_comparision_v2.png -------------------------------------------------------------------------------- /assets/speed_comparision_v3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/speed_comparision_v3.png -------------------------------------------------------------------------------- /assets/train_batch.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/train_batch.jpg -------------------------------------------------------------------------------- /assets/voc_loss_curve.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/voc_loss_curve.jpg -------------------------------------------------------------------------------- /assets/wechat_qrcode.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/wechat_qrcode.png -------------------------------------------------------------------------------- /assets/yolov5s.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/yolov5s.jpg -------------------------------------------------------------------------------- /assets/yolov6s.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/yolov6s.jpg -------------------------------------------------------------------------------- /assets/yoloxs.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/assets/yoloxs.jpg -------------------------------------------------------------------------------- /configs/base/README.md: -------------------------------------------------------------------------------- 1 | ## YOLOv6 base model 2 | 3 | English | [简体中文](./README_cn.md) 4 | 5 | ### Features 6 | 7 | - Use only regular convolution and Relu activation functions. 8 | 9 | - Apply CSP (1/2 channel dim) blocks in the network structure, except for Nano base model. 10 | 11 | Advantage: 12 | - Adopt a unified network structure and configuration, and the accuracy loss of the PTQ 8-bit quantization model is negligible. 13 | - Suitable for users who are just getting started or who need to apply, optimize and deploy an 8-bit quantization model quickly and frequently. 14 | 15 | 16 | ### Performance 17 | 18 | | Model | Size | mAPval
0.5:0.95 | SpeedT4
TRT FP16 b1
(FPS) | SpeedT4
TRT FP16 b32
(FPS) | SpeedT4
TRT INT8 b1
(FPS) | SpeedT4
TRT INT8 b32
(FPS) | Params
(M) | FLOPs
(G) | 19 | | :--------------------------------------------------------------------------------------------- | --- | ----------------- | ----- | ---- | ---- | ---- | ----- | ------ | 20 | | [**YOLOv6-N-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6n_base.pt) | 640 | 36.6distill | 727 | 1302 | 814 | 1805 | 4.65 | 11.46 | 21 | | [**YOLOv6-S-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6s_base.pt) | 640 | 45.3distill | 346 | 525 | 487 | 908 | 13.14 | 30.6 | 22 | | [**YOLOv6-M-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6m_base.pt) | 640 | 49.4distill | 179 | 245 | 284 | 439 | 28.33 | 72.30 | 23 | | [**YOLOv6-L-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6l_base.pt) | 640 | 51.1distill | 116 | 157 | 196 | 288 | 59.61 | 150.89 | 24 | 25 | - Speed is tested with TensorRT 8.2.4.2 on T4. 26 | - The processes of model training, evaluation, and inference are the same as the original ones. For details, please refer to [this README](https://github.com/meituan/YOLOv6#quick-start). -------------------------------------------------------------------------------- /configs/base/README_cn.md: -------------------------------------------------------------------------------- 1 | ## YOLOv6 基础版模型 2 | 3 | 简体中文 | [English](./README.md) 4 | 5 | ### 模型特点 6 | 7 | - 仅使用常规卷积和Relu激活函数 8 | 9 | - 网络结构均采用CSP (1/2通道) block，Nano网络除外。 10 | 11 | 优势： 12 | - 采用统一的网络结构和配置，且 PTQ 8位量化模型精度损失较小，适合刚入门或有快速迭代部署8位量化模型需求的用户。 13 | 14 | 15 | ### 模型指标 16 | 17 | | 模型 | 尺寸 | mAPval
0.5:0.95 | 速度T4
TRT FP16 b1
(FPS) | 速度T4
TRT FP16 b32
(FPS) | 速度T4
TRT INT8 b1
(FPS) | 速度T4
TRT INT8 b32
(FPS) | Params
(M) | FLOPs
(G) | 18 | | :--------------------------------------------------------------------------------------------- | --- | ----------------- | ----- | ---- | ---- | ---- | ----- | ------ | 19 | | [**YOLOv6-N-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6n_base.pt) | 640 | 36.6distill | 727 | 1302 | 814 | 1805 | 4.65 | 11.46 | 20 | | [**YOLOv6-S-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6s_base.pt) | 640 | 45.3distill | 346 | 525 | 487 | 908 | 13.14 | 30.6 | 21 | | [**YOLOv6-M-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6m_base.pt) | 640 | 49.4distill | 179 | 245 | 284 | 439 | 28.33 | 72.30 | 22 | | [**YOLOv6-L-base**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6l_base.pt) | 640 | 51.1distill | 116 | 157 | 196 | 288 | 59.61 | 150.89 | 23 | 24 | - 速度是在 T4 上测试的，TensorRT 版本为 8.4.2.4； 25 | - 模型训练、评估、推理流程与原来保持一致，具体可参考 [首页 README 文档](https://github.com/meituan/YOLOv6/blob/main/README_cn.md#%E5%BF%AB%E9%80%9F%E5%BC%80%E5%A7%8B)。 26 | -------------------------------------------------------------------------------- /configs/base/yolov6l_base.py: -------------------------------------------------------------------------------- 1 | # YOLOv6l large base model 2 | model = dict( 3 | type='YOLOv6l_base', 4 | pretrained=None, 5 | depth_multiple=1.0, 6 | width_multiple=1.0, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 2.0, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver=dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.01, 46 | lrf=0.01, 47 | momentum=0.937, 48 | weight_decay=0.0005, 49 | warmup_epochs=3.0, 50 | warmup_momentum=0.8, 51 | warmup_bias_lr=0.1 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.015, 56 | hsv_s=0.7, 57 | hsv_v=0.4, 58 | degrees=0.0, 59 | translate=0.1, 60 | scale=0.9, 61 | shear=0.0, 62 | flipud=0.0, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.1, 66 | ) 67 | training_mode = "conv_relu" 68 | -------------------------------------------------------------------------------- /configs/base/yolov6l_base_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6 large base model 2 | model = dict( 3 | type='YOLOv6l_base', 4 | depth_multiple=1.0, 5 | width_multiple=1.0, 6 | pretrained=None, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | ), 13 | neck=dict( 14 | type='CSPRepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | csp_e=float(1)/2, 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=1, 25 | out_indices=[17, 20, 23], 26 | strides=[8, 16, 32], 27 | iou_type='giou', 28 | use_dfl=True, 29 | reg_max=16, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 2.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.0032, 41 | lrf=0.12, 42 | momentum=0.843, 43 | weight_decay=0.00036, 44 | warmup_epochs=2.0, 45 | warmup_momentum=0.5, 46 | warmup_bias_lr=0.05 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.0138, 51 | hsv_s=0.664, 52 | hsv_v=0.464, 53 | degrees=0.373, 54 | translate=0.245, 55 | scale=0.898, 56 | shear=0.602, 57 | flipud=0.00856, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.243, 61 | ) 62 | training_mode = "conv_relu" 63 | -------------------------------------------------------------------------------- /configs/base/yolov6m_base.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m medium/large base model 2 | model = dict( 3 | type='YOLOv6m_base', 4 | pretrained=None, 5 | depth_multiple=0.80, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 0.8, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver=dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.01, 46 | lrf=0.01, 47 | momentum=0.937, 48 | weight_decay=0.0005, 49 | warmup_epochs=3.0, 50 | warmup_momentum=0.8, 51 | warmup_bias_lr=0.1 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.015, 56 | hsv_s=0.7, 57 | hsv_v=0.4, 58 | degrees=0.0, 59 | translate=0.1, 60 | scale=0.9, 61 | shear=0.0, 62 | flipud=0.0, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.1, 66 | ) 67 | training_mode = "conv_relu" -------------------------------------------------------------------------------- /configs/base/yolov6m_base_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m medium/large base model 2 | model = dict( 3 | type='YOLOv6m_base', 4 | pretrained=None, 5 | depth_multiple=0.80, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 0.8, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver = dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.0032, 46 | lrf=0.12, 47 | momentum=0.843, 48 | weight_decay=0.00036, 49 | warmup_epochs=2.0, 50 | warmup_momentum=0.5, 51 | warmup_bias_lr=0.05 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.0138, 56 | hsv_s=0.664, 57 | hsv_v=0.464, 58 | degrees=0.373, 59 | translate=0.245, 60 | scale=0.898, 61 | shear=0.602, 62 | flipud=0.00856, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.243, 66 | ) 67 | training_mode = "conv_relu" 68 | -------------------------------------------------------------------------------- /configs/base/yolov6n_base.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s nano base model 2 | model = dict( 3 | type='YOLOv6n_base', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='giou', 32 | use_dfl=True, # set to True if you want to further train with distillation 33 | reg_max=16, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.01, 45 | lrf=0.01, 46 | momentum=0.937, 47 | weight_decay=0.0005, 48 | warmup_epochs=3.0, 49 | warmup_momentum=0.8, 50 | warmup_bias_lr=0.1 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.015, 55 | hsv_s=0.7, 56 | hsv_v=0.4, 57 | degrees=0.0, 58 | translate=0.1, 59 | scale=0.5, 60 | shear=0.0, 61 | flipud=0.0, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.0, 65 | ) 66 | training_mode = "conv_relu" -------------------------------------------------------------------------------- /configs/base/yolov6n_base_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s nanao base model 2 | model = dict( 3 | type='YOLOv6n_base', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='giou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.0032, 45 | lrf=0.12, 46 | momentum=0.843, 47 | weight_decay=0.00036, 48 | warmup_epochs=2.0, 49 | warmup_momentum=0.5, 50 | warmup_bias_lr=0.05 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.0138, 55 | hsv_s=0.664, 56 | hsv_v=0.464, 57 | degrees=0.373, 58 | translate=0.245, 59 | scale=0.898, 60 | shear=0.602, 61 | flipud=0.00856, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.243, 65 | ) 66 | training_mode = "conv_relu" 67 | -------------------------------------------------------------------------------- /configs/base/yolov6s_base.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s small base model 2 | model = dict( 3 | type='YOLOv6s_base', 4 | pretrained=None, 5 | depth_multiple=0.70, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | cspsppf=True, 14 | ), 15 | neck=dict( 16 | type='CSPRepBiFPANNeck',#CSPRepPANNeck 17 | num_repeats=[12, 12, 12, 12], 18 | out_channels=[256, 128, 128, 256, 256, 512], 19 | csp_e=float(1)/2, 20 | ), 21 | head=dict( 22 | type='EffiDeHead', 23 | in_channels=[128, 256, 512], 24 | num_layers=3, 25 | begin_indices=24, 26 | anchors=3, 27 | anchors_init=[[10,13, 19,19, 33,23], 28 | [30,61, 59,59, 59,119], 29 | [116,90, 185,185, 373,326]], 30 | out_indices=[17, 20, 23], 31 | strides=[8, 16, 32], 32 | atss_warmup_epoch=0, 33 | iou_type='giou', 34 | use_dfl=True, # set to True if you want to further train with distillation 35 | reg_max=16, # set to 16 if you want to further train with distillation 36 | distill_weight={ 37 | 'class': 1.0, 38 | 'dfl': 1.0, 39 | }, 40 | ) 41 | ) 42 | 43 | solver = dict( 44 | optim='SGD', 45 | lr_scheduler='Cosine', 46 | lr0=0.01, 47 | lrf=0.01, 48 | momentum=0.937, 49 | weight_decay=0.0005, 50 | warmup_epochs=3.0, 51 | warmup_momentum=0.8, 52 | warmup_bias_lr=0.1 53 | ) 54 | 55 | data_aug = dict( 56 | hsv_h=0.015, 57 | hsv_s=0.7, 58 | hsv_v=0.4, 59 | degrees=0.0, 60 | translate=0.1, 61 | scale=0.5, 62 | shear=0.0, 63 | flipud=0.0, 64 | fliplr=0.5, 65 | mosaic=1.0, 66 | mixup=0.0, 67 | ) 68 | training_mode = "conv_relu" 69 | -------------------------------------------------------------------------------- /configs/base/yolov6s_base_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s small base model 2 | model = dict( 3 | type='YOLOv6s_base', 4 | pretrained=None, 5 | depth_multiple=0.70, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | cspsppf=True, 14 | ), 15 | neck=dict( 16 | type='CSPRepPANNeck', 17 | num_repeats=[12, 12, 12, 12], 18 | out_channels=[256, 128, 128, 256, 256, 512], 19 | csp_e=float(1)/2, 20 | ), 21 | head=dict( 22 | type='EffiDeHead', 23 | in_channels=[128, 256, 512], 24 | num_layers=3, 25 | begin_indices=24, 26 | anchors=3, 27 | anchors_init=[[10,13, 19,19, 33,23], 28 | [30,61, 59,59, 59,119], 29 | [116,90, 185,185, 373,326]], 30 | out_indices=[17, 20, 23], 31 | strides=[8, 16, 32], 32 | atss_warmup_epoch=0, 33 | iou_type='giou', 34 | use_dfl=False, # set to True if you want to further train with distillation 35 | reg_max=0, # set to 16 if you want to further train with distillation 36 | distill_weight={ 37 | 'class': 1.0, 38 | 'dfl': 1.0, 39 | }, 40 | ) 41 | ) 42 | 43 | solver = dict( 44 | optim='SGD', 45 | lr_scheduler='Cosine', 46 | lr0=0.0032, 47 | lrf=0.12, 48 | momentum=0.843, 49 | weight_decay=0.00036, 50 | warmup_epochs=2.0, 51 | warmup_momentum=0.5, 52 | warmup_bias_lr=0.05 53 | ) 54 | 55 | data_aug = dict( 56 | hsv_h=0.0138, 57 | hsv_s=0.664, 58 | hsv_v=0.464, 59 | degrees=0.373, 60 | translate=0.245, 61 | scale=0.898, 62 | shear=0.602, 63 | flipud=0.00856, 64 | fliplr=0.5, 65 | mosaic=1.0, 66 | mixup=0.243, 67 | ) 68 | training_mode = "conv_relu" 69 | -------------------------------------------------------------------------------- /configs/experiment/eval_640_repro.py: -------------------------------------------------------------------------------- 1 | # eval param for different scale 2 | 3 | eval_params = dict( 4 | default = dict( 5 | img_size=640, 6 | shrink_size=2, 7 | infer_on_rect=False, 8 | ), 9 | yolov6n = dict( 10 | img_size=640, 11 | shrink_size=4, 12 | infer_on_rect=False, 13 | ), 14 | yolov6t = dict( 15 | img_size=640, 16 | shrink_size=6, 17 | infer_on_rect=False, 18 | ), 19 | yolov6s = dict( 20 | img_size=640, 21 | shrink_size=6, 22 | infer_on_rect=False, 23 | ), 24 | yolov6m = dict( 25 | img_size=640, 26 | shrink_size=4, 27 | infer_on_rect=False, 28 | ), 29 | yolov6l = dict( 30 | img_size=640, 31 | shrink_size=4, 32 | infer_on_rect=False, 33 | ), 34 | yolov6l_relu = dict( 35 | img_size=640, 36 | shrink_size=2, 37 | infer_on_rect=False, 38 | ), 39 | yolov6n6 = dict( 40 | img_size=1280, 41 | shrink_size=17, 42 | infer_on_rect=False, 43 | ), 44 | yolov6s6 = dict( 45 | img_size=1280, 46 | shrink_size=8, 47 | infer_on_rect=False, 48 | ), 49 | yolov6m6 = dict( 50 | img_size=1280, 51 | shrink_size=64, 52 | infer_on_rect=False, 53 | ), 54 | yolov6l6 = dict( 55 | img_size=1280, 56 | shrink_size=41, 57 | infer_on_rect=False, 58 | ) 59 | ) 60 | -------------------------------------------------------------------------------- /configs/experiment/yolov6n_with_eval_params.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model with eval param(when traing) 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | ), 12 | neck=dict( 13 | type='RepPANNeck', 14 | num_repeats=[12, 12, 12, 12], 15 | out_channels=[256, 128, 128, 256, 256, 512], 16 | ), 17 | head=dict( 18 | type='EffiDeHead', 19 | in_channels=[128, 256, 512], 20 | num_layers=3, 21 | begin_indices=24, 22 | anchors=1, 23 | out_indices=[17, 20, 23], 24 | strides=[8, 16, 32], 25 | iou_type='siou', 26 | use_dfl=False, 27 | reg_max=0 #if use_dfl is False, please set reg_max to 0 28 | ) 29 | ) 30 | 31 | solver = dict( 32 | optim='SGD', 33 | lr_scheduler='Cosine', 34 | lr0=0.02, #0.01 # 0.02 35 | lrf=0.01, 36 | momentum=0.937, 37 | weight_decay=0.0005, 38 | warmup_epochs=3.0, 39 | warmup_momentum=0.8, 40 | warmup_bias_lr=0.1 41 | ) 42 | 43 | data_aug = dict( 44 | hsv_h=0.015, 45 | hsv_s=0.7, 46 | hsv_v=0.4, 47 | degrees=0.0, 48 | translate=0.1, 49 | scale=0.5, 50 | shear=0.0, 51 | flipud=0.0, 52 | fliplr=0.5, 53 | mosaic=1.0, 54 | mixup=0.0, 55 | ) 56 | 57 | # Eval params when eval model. 58 | # If eval_params item is list, eg conf_thres=[0.03, 0.03], 59 | # first will be used in train.py and second will be used in eval.py. 60 | eval_params = dict( 61 | batch_size=None, #None mean will be the same as batch on one device * 2 62 | img_size=None, #None mean will be the same as train image size 63 | conf_thres=0.03, 64 | iou_thres=0.65, 65 | 66 | #pading and scale coord 67 | shrink_size=None, # None mean will not shrink the image. 68 | infer_on_rect=True, 69 | 70 | #metric 71 | verbose=False, 72 | do_coco_metric=True, 73 | do_pr_metric=False, 74 | plot_curve=False, 75 | plot_confusion_matrix=False 76 | ) 77 | -------------------------------------------------------------------------------- /configs/experiment/yolov6s_csp_scaled.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m model 2 | model = dict( 3 | type='YOLOv6s_csp', 4 | pretrained=None, 5 | depth_multiple=0.70, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | ), 13 | neck=dict( 14 | type='CSPRepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | csp_e=float(1)/2, 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=1, 25 | out_indices=[17, 20, 23], 26 | strides=[8, 16, 32], 27 | iou_type='giou', 28 | use_dfl=False, 29 | reg_max=0 #if use_dfl is False, please set reg_max to 0 30 | ) 31 | ) 32 | 33 | solver=dict( 34 | optim='SGD', 35 | lr_scheduler='Cosine', 36 | lr0=0.01, 37 | lrf=0.01, 38 | momentum=0.937, 39 | weight_decay=0.0005, 40 | warmup_epochs=3.0, 41 | warmup_momentum=0.8, 42 | warmup_bias_lr=0.1 43 | ) 44 | 45 | data_aug = dict( 46 | hsv_h=0.015, 47 | hsv_s=0.7, 48 | hsv_v=0.4, 49 | degrees=0.0, 50 | translate=0.1, 51 | scale=0.9, 52 | shear=0.0, 53 | flipud=0.0, 54 | fliplr=0.5, 55 | mosaic=1.0, 56 | mixup=0.1, 57 | ) 58 | -------------------------------------------------------------------------------- /configs/experiment/yolov6t.py: -------------------------------------------------------------------------------- 1 | # YOLOv6t model 2 | model = dict( 3 | type='YOLOv6t', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.375, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | ), 12 | neck=dict( 13 | type='RepPANNeck', 14 | num_repeats=[12, 12, 12, 12], 15 | out_channels=[256, 128, 128, 256, 256, 512], 16 | ), 17 | head=dict( 18 | type='EffiDeHead', 19 | in_channels=[128, 256, 512], 20 | num_layers=3, 21 | begin_indices=24, 22 | anchors=1, 23 | out_indices=[17, 20, 23], 24 | strides=[8, 16, 32], 25 | iou_type='siou', 26 | use_dfl=False, 27 | reg_max=0 #if use_dfl is False, please set reg_max to 0 28 | ) 29 | ) 30 | 31 | solver = dict( 32 | optim='SGD', 33 | lr_scheduler='Cosine', 34 | lr0=0.01, 35 | lrf=0.01, 36 | momentum=0.937, 37 | weight_decay=0.0005, 38 | warmup_epochs=3.0, 39 | warmup_momentum=0.8, 40 | warmup_bias_lr=0.1 41 | ) 42 | 43 | data_aug = dict( 44 | hsv_h=0.015, 45 | hsv_s=0.7, 46 | hsv_v=0.4, 47 | degrees=0.0, 48 | translate=0.1, 49 | scale=0.5, 50 | shear=0.0, 51 | flipud=0.0, 52 | fliplr=0.5, 53 | mosaic=1.0, 54 | mixup=0.0, 55 | ) 56 | -------------------------------------------------------------------------------- /configs/experiment/yolov6t_csp_scaled.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n_csp', 4 | pretrained=None, 5 | depth_multiple=0.60, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | ), 13 | neck=dict( 14 | type='CSPRepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | csp_e=float(1)/2, 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=1, 25 | out_indices=[17, 20, 23], 26 | strides=[8, 16, 32], 27 | iou_type='giou', 28 | use_dfl=False, 29 | reg_max=0 #if use_dfl is False, please set reg_max to 0 30 | ) 31 | ) 32 | 33 | solver=dict( 34 | optim='SGD', 35 | lr_scheduler='Cosine', 36 | lr0=0.01, 37 | lrf=0.01, 38 | momentum=0.937, 39 | weight_decay=0.0005, 40 | warmup_epochs=3.0, 41 | warmup_momentum=0.8, 42 | warmup_bias_lr=0.1 43 | ) 44 | 45 | data_aug = dict( 46 | hsv_h=0.015, 47 | hsv_s=0.7, 48 | hsv_v=0.4, 49 | degrees=0.0, 50 | translate=0.1, 51 | scale=0.9, 52 | shear=0.0, 53 | flipud=0.0, 54 | fliplr=0.5, 55 | mosaic=1.0, 56 | mixup=0.1, 57 | ) 58 | -------------------------------------------------------------------------------- /configs/experiment/yolov6t_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6t model 2 | model = dict( 3 | type='YOLOv6t', 4 | pretrained='weights/yolov6t.pt', 5 | depth_multiple=0.33, 6 | width_multiple=0.375, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | ), 12 | neck=dict( 13 | type='RepPANNeck', 14 | num_repeats=[12, 12, 12, 12], 15 | out_channels=[256, 128, 128, 256, 256, 512], 16 | ), 17 | head=dict( 18 | type='EffiDeHead', 19 | in_channels=[128, 256, 512], 20 | num_layers=3, 21 | begin_indices=24, 22 | anchors=1, 23 | out_indices=[17, 20, 23], 24 | strides=[8, 16, 32], 25 | iou_type='siou', 26 | use_dfl=False, 27 | reg_max=0 #if use_dfl is False, please set reg_max to 0 28 | ) 29 | ) 30 | 31 | solver = dict( 32 | optim='SGD', 33 | lr_scheduler='Cosine', 34 | lr0=0.0032, 35 | lrf=0.12, 36 | momentum=0.843, 37 | weight_decay=0.00036, 38 | warmup_epochs=2.0, 39 | warmup_momentum=0.5, 40 | warmup_bias_lr=0.05 41 | ) 42 | 43 | data_aug = dict( 44 | hsv_h=0.0138, 45 | hsv_s=0.664, 46 | hsv_v=0.464, 47 | degrees=0.373, 48 | translate=0.245, 49 | scale=0.898, 50 | shear=0.602, 51 | flipud=0.00856, 52 | fliplr=0.5, 53 | mosaic=1.0, 54 | mixup=0.243, 55 | ) 56 | -------------------------------------------------------------------------------- /configs/qarepvgg/README.md: -------------------------------------------------------------------------------- 1 | ## YOLOv6 base model 2 | 3 | English | [简体中文](./README_cn.md) 4 | 5 | ### Features 6 | 7 | - This is a RepOpt-version implementation of YOLOv6 according to [QARepVGG](https://arxiv.org/abs/2212.01593). 8 | 9 | - The QARep version models possess slightly lower float accuracy on COCO than the RepVGG version models, but achieve highly improved quantized accuracy. 10 | 11 | - The INT8 accuracies listed were obtained using a simple PTQ process, as implemented in the [`onnx_to_trt.py`](../../deploy/TensorRT/onnx_to_trt.py) script. However, higher accuracies could be achieved using Quantization-Aware Training (QAT) due to the specific architecture design of the QARepVGG model. 12 | 13 | ### Performance 14 | 15 | | Model | Size | Float
mAPval
0.5:0.95 | INT8
mAPval
0.5:0.95 | SpeedT4
trt fp16 b32
(fps) | SpeedT4
trt int8 b32
(fps) | Params
(M) | FLOPs
(G) | 16 | | :----------------------------------------------------------- | -------- | :----------------------- | -------------------------------------- | --------------------------------------- | -------------------- | ------------------- | -------------------- | 17 | | [**YOLOv6-N**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6n.pt) | 640 | 37.5 | 34.3 | 1286 | 1773 |4.7 | 11.4 | 18 | | [**YOLOv6-N-qa**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6n_qa.pt) | 640 | 37.1 | 36.4 | 1286 | 1773 | 4.7 | 11.4 | 19 | | [**YOLOv6-S**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6s.pt) | 640 | 45.0 | 41.3 | 513 | 1117 | 18.5 | 45.3 | 20 | | [**YOLOv6-S-qa**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6s_qa.pt) | 640 | 44.7 | 44.0 | 513 | 1117 | 18.5 | 45.3 | 21 | | [**YOLOv6-M**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6m.pt) | 640 | 50.0 | 48.1 | 250 | 439 | 34.9 | 85.8 | 22 | | [**YOLOv6-M-qa**](https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6m_qa.pt) | 640 | 49.7 | 49.4 | 250 | 439 | 34.9 | 85.8 | 23 | 24 | - Speed is tested with TensorRT 8.4 on T4. 25 | - We have not conducted experiments on the YOLOv6-L model since it does not use the RepVGG architecture. 26 | - The processes of model training, evaluation, and inference are the same as the original ones. For details, please refer to [this README](https://github.com/meituan/YOLOv6#quick-start). 27 | -------------------------------------------------------------------------------- /configs/qarepvgg/yolov6m_qa.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m model 2 | model = dict( 3 | type='YOLOv6m', 4 | pretrained=None, 5 | depth_multiple=0.60, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(2)/3, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(2)/3, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 0.8, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver=dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.01, 46 | lrf=0.01, 47 | momentum=0.937, 48 | weight_decay=0.0005, 49 | warmup_epochs=3.0, 50 | warmup_momentum=0.8, 51 | warmup_bias_lr=0.1 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.015, 56 | hsv_s=0.7, 57 | hsv_v=0.4, 58 | degrees=0.0, 59 | translate=0.1, 60 | scale=0.9, 61 | shear=0.0, 62 | flipud=0.0, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.1, 66 | ) 67 | 68 | training_mode='qarepvggv2' -------------------------------------------------------------------------------- /configs/qarepvgg/yolov6n_qa.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='siou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.02, 45 | lrf=0.01, 46 | momentum=0.937, 47 | weight_decay=0.0005, 48 | warmup_epochs=3.0, 49 | warmup_momentum=0.8, 50 | warmup_bias_lr=0.1 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.015, 55 | hsv_s=0.7, 56 | hsv_v=0.4, 57 | degrees=0.0, 58 | translate=0.1, 59 | scale=0.5, 60 | shear=0.0, 61 | flipud=0.0, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.0, 65 | ) 66 | training_mode='qarepvggv2' 67 | -------------------------------------------------------------------------------- /configs/qarepvgg/yolov6s_qa.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6s', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='giou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.01, 45 | lrf=0.01, 46 | momentum=0.937, 47 | weight_decay=0.0005, 48 | warmup_epochs=3.0, 49 | warmup_momentum=0.8, 50 | warmup_bias_lr=0.1 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.015, 55 | hsv_s=0.7, 56 | hsv_v=0.4, 57 | degrees=0.0, 58 | translate=0.1, 59 | scale=0.5, 60 | shear=0.0, 61 | flipud=0.0, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.0, 65 | ) 66 | 67 | training_mode='qarepvggv2' -------------------------------------------------------------------------------- /configs/repopt/yolov6_tiny_hs.py: -------------------------------------------------------------------------------- 1 | # YOLOv6t model 2 | model = dict( 3 | type='YOLOv6t', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.375, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | ), 12 | neck=dict( 13 | type='RepPANNeck', 14 | num_repeats=[12, 12, 12, 12], 15 | out_channels=[256, 128, 128, 256, 256, 512], 16 | ), 17 | head=dict( 18 | type='EffiDeHead', 19 | in_channels=[128, 256, 512], 20 | num_layers=3, 21 | begin_indices=24, 22 | anchors=1, 23 | out_indices=[17, 20, 23], 24 | strides=[8, 16, 32], 25 | atss_warmup_epoch=0, 26 | iou_type='siou', 27 | use_dfl=False, 28 | reg_max=0 #if use_dfl is False, please set reg_max to 0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.01, 36 | lrf=0.01, 37 | momentum=0.937, 38 | weight_decay=0.0005, 39 | warmup_epochs=3.0, 40 | warmup_momentum=0.8, 41 | warmup_bias_lr=0.1 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.015, 46 | hsv_s=0.7, 47 | hsv_v=0.4, 48 | degrees=0.0, 49 | translate=0.1, 50 | scale=0.5, 51 | shear=0.0, 52 | flipud=0.0, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.0, 56 | ) 57 | 58 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 59 | training_mode='hyper_search' 60 | -------------------------------------------------------------------------------- /configs/repopt/yolov6_tiny_opt.py: -------------------------------------------------------------------------------- 1 | # YOLOv6t model 2 | model = dict( 3 | type='YOLOv6t', 4 | pretrained=None, 5 | scales='../yolov6_assert/v6t_v2_scale_last.pt', 6 | depth_multiple=0.33, 7 | width_multiple=0.375, 8 | backbone=dict( 9 | type='EfficientRep', 10 | num_repeats=[1, 6, 12, 18, 6], 11 | out_channels=[64, 128, 256, 512, 1024], 12 | ), 13 | neck=dict( 14 | type='RepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | ), 18 | head=dict( 19 | type='EffiDeHead', 20 | in_channels=[128, 256, 512], 21 | num_layers=3, 22 | begin_indices=24, 23 | anchors=1, 24 | out_indices=[17, 20, 23], 25 | strides=[8, 16, 32], 26 | atss_warmup_epoch=0, 27 | iou_type='siou', 28 | use_dfl=False, 29 | reg_max=0 #if use_dfl is False, please set reg_max to 0 30 | ) 31 | ) 32 | 33 | solver = dict( 34 | optim='SGD', 35 | lr_scheduler='Cosine', 36 | lr0=0.01, 37 | lrf=0.01, 38 | momentum=0.937, 39 | weight_decay=0.0005, 40 | warmup_epochs=3.0, 41 | warmup_momentum=0.8, 42 | warmup_bias_lr=0.1 43 | ) 44 | 45 | data_aug = dict( 46 | hsv_h=0.015, 47 | hsv_s=0.7, 48 | hsv_v=0.4, 49 | degrees=0.0, 50 | translate=0.1, 51 | scale=0.5, 52 | shear=0.0, 53 | flipud=0.0, 54 | fliplr=0.5, 55 | mosaic=1.0, 56 | mixup=0.0, 57 | ) 58 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 59 | training_mode='repopt' 60 | -------------------------------------------------------------------------------- /configs/repopt/yolov6_tiny_opt_qat.py: -------------------------------------------------------------------------------- 1 | # YOLOv6t model 2 | model = dict( 3 | type='YOLOv6t', 4 | pretrained='./assets/v6s_t.pt', 5 | scales='./assets/v6t_v2_scale_last.pt', 6 | depth_multiple=0.33, 7 | width_multiple=0.375, 8 | backbone=dict( 9 | type='EfficientRep', 10 | num_repeats=[1, 6, 12, 18, 6], 11 | out_channels=[64, 128, 256, 512, 1024], 12 | ), 13 | neck=dict( 14 | type='RepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | ), 18 | head=dict( 19 | type='EffiDeHead', 20 | in_channels=[128, 256, 512], 21 | num_layers=3, 22 | begin_indices=24, 23 | anchors=1, 24 | out_indices=[17, 20, 23], 25 | strides=[8, 16, 32], 26 | atss_warmup_epoch=0, 27 | iou_type='siou', 28 | use_dfl=False, 29 | reg_max=0, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.00001, 41 | lrf=0.001, 42 | momentum=0.937, 43 | weight_decay=0.00005, 44 | warmup_epochs=3.0, 45 | warmup_momentum=0.8, 46 | warmup_bias_lr=0.1 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.015, 51 | hsv_s=0.7, 52 | hsv_v=0.4, 53 | degrees=0.0, 54 | translate=0.1, 55 | scale=0.5, 56 | shear=0.0, 57 | flipud=0.0, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.0, 61 | ) 62 | 63 | ptq = dict( 64 | num_bits = 8, 65 | calib_batches = 4, 66 | # 'max', 'histogram' 67 | calib_method = 'max', 68 | # 'entropy', 'percentile', 'mse' 69 | histogram_amax_method='entropy', 70 | histogram_amax_percentile=99.99, 71 | calib_output_path='./', 72 | sensitive_layers_skip=False, 73 | sensitive_layers_list=[], 74 | ) 75 | 76 | qat = dict( 77 | calib_pt = './assets/v6s_t_calib_max.pt', 78 | sensitive_layers_skip = False, 79 | sensitive_layers_list=[], 80 | ) 81 | 82 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 83 | training_mode='repopt' 84 | -------------------------------------------------------------------------------- /configs/repopt/yolov6n_hs.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | ), 12 | neck=dict( 13 | type='RepPANNeck', 14 | num_repeats=[12, 12, 12, 12], 15 | out_channels=[256, 128, 128, 256, 256, 512], 16 | ), 17 | head=dict( 18 | type='EffiDeHead', 19 | in_channels=[128, 256, 512], 20 | num_layers=3, 21 | begin_indices=24, 22 | anchors=1, 23 | out_indices=[17, 20, 23], 24 | strides=[8, 16, 32], 25 | atss_warmup_epoch=0, 26 | iou_type='siou', 27 | use_dfl=False, 28 | reg_max=0 #if use_dfl is False, please set reg_max to 0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.02, #0.01 # 0.02 36 | lrf=0.01, 37 | momentum=0.937, 38 | weight_decay=0.0005, 39 | warmup_epochs=3.0, 40 | warmup_momentum=0.8, 41 | warmup_bias_lr=0.1 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.015, 46 | hsv_s=0.7, 47 | hsv_v=0.4, 48 | degrees=0.0, 49 | translate=0.1, 50 | scale=0.5, 51 | shear=0.0, 52 | flipud=0.0, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.0, 56 | ) 57 | 58 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 59 | training_mode='hyper_search' 60 | -------------------------------------------------------------------------------- /configs/repopt/yolov6n_opt.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained=None, 5 | scales='../yolov6_assert/v6n_v2_scale_last.pt', 6 | depth_multiple=0.33, 7 | width_multiple=0.25, 8 | backbone=dict( 9 | type='EfficientRep', 10 | num_repeats=[1, 6, 12, 18, 6], 11 | out_channels=[64, 128, 256, 512, 1024], 12 | ), 13 | neck=dict( 14 | type='RepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | ), 18 | head=dict( 19 | type='EffiDeHead', 20 | in_channels=[128, 256, 512], 21 | num_layers=3, 22 | begin_indices=24, 23 | anchors=1, 24 | out_indices=[17, 20, 23], 25 | strides=[8, 16, 32], 26 | atss_warmup_epoch=0, 27 | iou_type='siou', 28 | use_dfl=False, 29 | reg_max=0 #if use_dfl is False, please set reg_max to 0 30 | ) 31 | ) 32 | 33 | solver = dict( 34 | optim='SGD', 35 | lr_scheduler='Cosine', 36 | lr0=0.02, #0.01 # 0.02 37 | lrf=0.01, 38 | momentum=0.937, 39 | weight_decay=0.0005, 40 | warmup_epochs=3.0, 41 | warmup_momentum=0.8, 42 | warmup_bias_lr=0.1 43 | ) 44 | 45 | data_aug = dict( 46 | hsv_h=0.015, 47 | hsv_s=0.7, 48 | hsv_v=0.4, 49 | degrees=0.0, 50 | translate=0.1, 51 | scale=0.5, 52 | shear=0.0, 53 | flipud=0.0, 54 | fliplr=0.5, 55 | mosaic=1.0, 56 | mixup=0.0, 57 | ) 58 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 59 | training_mode='repopt' 60 | -------------------------------------------------------------------------------- /configs/repopt/yolov6n_opt_qat.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained='./assets/v6s_n.pt', 5 | scales='./assets/v6n_v2_scale_last.pt', 6 | depth_multiple=0.33, 7 | width_multiple=0.25, 8 | backbone=dict( 9 | type='EfficientRep', 10 | num_repeats=[1, 6, 12, 18, 6], 11 | out_channels=[64, 128, 256, 512, 1024], 12 | ), 13 | neck=dict( 14 | type='RepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | ), 18 | head=dict( 19 | type='EffiDeHead', 20 | in_channels=[128, 256, 512], 21 | num_layers=3, 22 | begin_indices=24, 23 | anchors=1, 24 | out_indices=[17, 20, 23], 25 | strides=[8, 16, 32], 26 | atss_warmup_epoch=0, 27 | iou_type='siou', 28 | use_dfl=False, 29 | reg_max=0, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.00001, #0.01 # 0.02 41 | lrf=0.001, 42 | momentum=0.937, 43 | weight_decay=0.00005, 44 | warmup_epochs=3.0, 45 | warmup_momentum=0.8, 46 | warmup_bias_lr=0.1 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.015, 51 | hsv_s=0.7, 52 | hsv_v=0.4, 53 | degrees=0.0, 54 | translate=0.1, 55 | scale=0.5, 56 | shear=0.0, 57 | flipud=0.0, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.0, 61 | ) 62 | 63 | ptq = dict( 64 | num_bits = 8, 65 | calib_batches = 4, 66 | # 'max', 'histogram' 67 | calib_method = 'max', 68 | # 'entropy', 'percentile', 'mse' 69 | histogram_amax_method='entropy', 70 | histogram_amax_percentile=99.99, 71 | calib_output_path='./', 72 | sensitive_layers_skip=False, 73 | sensitive_layers_list=[], 74 | ) 75 | 76 | qat = dict( 77 | calib_pt = './assets/v6s_n_calib_max.pt', 78 | sensitive_layers_skip = False, 79 | sensitive_layers_list=[], 80 | ) 81 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 82 | training_mode='repopt' 83 | -------------------------------------------------------------------------------- /configs/repopt/yolov6s_hs.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6s', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | ), 12 | neck=dict( 13 | type='RepPANNeck', 14 | num_repeats=[12, 12, 12, 12], 15 | out_channels=[256, 128, 128, 256, 256, 512], 16 | ), 17 | head=dict( 18 | type='EffiDeHead', 19 | in_channels=[128, 256, 512], 20 | num_layers=3, 21 | begin_indices=24, 22 | anchors=1, 23 | out_indices=[17, 20, 23], 24 | strides=[8, 16, 32], 25 | atss_warmup_epoch=0, 26 | iou_type='giou', 27 | use_dfl=False, 28 | reg_max=0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.01, 36 | lrf=0.01, 37 | momentum=0.937, 38 | weight_decay=0.0005, 39 | warmup_epochs=3.0, 40 | warmup_momentum=0.8, 41 | warmup_bias_lr=0.1 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.015, 46 | hsv_s=0.7, 47 | hsv_v=0.4, 48 | degrees=0.0, 49 | translate=0.1, 50 | scale=0.5, 51 | shear=0.0, 52 | flipud=0.0, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.0, 56 | ) 57 | 58 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 59 | training_mode='hyper_search' 60 | -------------------------------------------------------------------------------- /configs/repopt/yolov6s_opt.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6s', 4 | pretrained=None, 5 | scales='../yolov6_assert/v6s_v2_scale.pt', 6 | depth_multiple=0.33, 7 | width_multiple=0.50, 8 | backbone=dict( 9 | type='EfficientRep', 10 | num_repeats=[1, 6, 12, 18, 6], 11 | out_channels=[64, 128, 256, 512, 1024], 12 | ), 13 | neck=dict( 14 | type='RepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | ), 18 | head=dict( 19 | type='EffiDeHead', 20 | in_channels=[128, 256, 512], 21 | num_layers=3, 22 | begin_indices=24, 23 | anchors=1, 24 | out_indices=[17, 20, 23], 25 | strides=[8, 16, 32], 26 | atss_warmup_epoch=0, 27 | iou_type='giou', 28 | use_dfl=False, 29 | reg_max=0 30 | ) 31 | ) 32 | 33 | solver = dict( 34 | optim='SGD', 35 | lr_scheduler='Cosine', 36 | lr0=0.01, 37 | lrf=0.01, 38 | momentum=0.937, 39 | weight_decay=0.0005, 40 | warmup_epochs=3.0, 41 | warmup_momentum=0.8, 42 | warmup_bias_lr=0.1 43 | ) 44 | 45 | data_aug = dict( 46 | hsv_h=0.015, 47 | hsv_s=0.7, 48 | hsv_v=0.4, 49 | degrees=0.0, 50 | translate=0.1, 51 | scale=0.5, 52 | shear=0.0, 53 | flipud=0.0, 54 | fliplr=0.5, 55 | mosaic=1.0, 56 | mixup=0.0, 57 | ) 58 | 59 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 60 | training_mode='repopt' 61 | -------------------------------------------------------------------------------- /configs/repopt/yolov6s_opt_qat.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6s', 4 | pretrained='./assets/yolov6s_v2_reopt_43.1.pt', 5 | scales='./assets/yolov6s_v2_scale.pt', 6 | depth_multiple=0.33, 7 | width_multiple=0.50, 8 | backbone=dict( 9 | type='EfficientRep', 10 | num_repeats=[1, 6, 12, 18, 6], 11 | out_channels=[64, 128, 256, 512, 1024], 12 | ), 13 | neck=dict( 14 | type='RepPANNeck', 15 | num_repeats=[12, 12, 12, 12], 16 | out_channels=[256, 128, 128, 256, 256, 512], 17 | ), 18 | head=dict( 19 | type='EffiDeHead', 20 | in_channels=[128, 256, 512], 21 | num_layers=3, 22 | begin_indices=24, 23 | anchors=1, 24 | out_indices=[17, 20, 23], 25 | strides=[8, 16, 32], 26 | atss_warmup_epoch=0, 27 | iou_type = 'giou', 28 | use_dfl = False, 29 | reg_max = 0, # if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.00001, 41 | lrf=0.001, 42 | momentum=0.937, 43 | weight_decay=0.00005, 44 | warmup_epochs=3, 45 | warmup_momentum=0.8, 46 | warmup_bias_lr=0.1 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.015, 51 | hsv_s=0.7, 52 | hsv_v=0.4, 53 | degrees=0.0, 54 | translate=0.1, 55 | scale=0.5, 56 | shear=0.0, 57 | flipud=0.0, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.0, 61 | ) 62 | 63 | ptq = dict( 64 | num_bits = 8, 65 | calib_batches = 4, 66 | # 'max', 'histogram' 67 | calib_method = 'histogram', 68 | # 'entropy', 'percentile', 'mse' 69 | histogram_amax_method='entropy', 70 | histogram_amax_percentile=99.99, 71 | calib_output_path='./', 72 | sensitive_layers_skip=False, 73 | sensitive_layers_list=['detect.stems.0.conv', 74 | 'detect.stems.1.conv', 75 | 'detect.stems.2.conv', 76 | 'detect.cls_convs.0.conv', 77 | 'detect.cls_convs.1.conv', 78 | 'detect.cls_convs.2.conv', 79 | 'detect.reg_convs.0.conv', 80 | 'detect.reg_convs.1.conv', 81 | 'detect.reg_convs.2.conv', 82 | 'detect.cls_preds.0', 83 | 'detect.cls_preds.1', 84 | 'detect.cls_preds.2', 85 | 'detect.reg_preds.0', 86 | 'detect.reg_preds.1', 87 | 'detect.reg_preds.2', 88 | ], 89 | ) 90 | 91 | qat = dict( 92 | calib_pt = './assets/yolov6s_v2_reopt_43.1_calib_histogram.pt', 93 | sensitive_layers_skip = False, 94 | sensitive_layers_list=['detect.stems.0.conv', 95 | 'detect.stems.1.conv', 96 | 'detect.stems.2.conv', 97 | 'detect.cls_convs.0.conv', 98 | 'detect.cls_convs.1.conv', 99 | 'detect.cls_convs.2.conv', 100 | 'detect.reg_convs.0.conv', 101 | 'detect.reg_convs.1.conv', 102 | 'detect.reg_convs.2.conv', 103 | 'detect.cls_preds.0', 104 | 'detect.cls_preds.1', 105 | 'detect.cls_preds.2', 106 | 'detect.reg_preds.0', 107 | 'detect.reg_preds.1', 108 | 'detect.reg_preds.2', 109 | ], 110 | ) 111 | 112 | # Choose Rep-block by the training Mode, choices=["repvgg", "hyper-search", "repopt"] 113 | training_mode='repopt' 114 | -------------------------------------------------------------------------------- /configs/yolov6l.py: -------------------------------------------------------------------------------- 1 | # YOLOv6l model 2 | model = dict( 3 | type='YOLOv6l', 4 | pretrained=None, 5 | depth_multiple=1.0, 6 | width_multiple=1.0, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 2.0, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver=dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.01, 46 | lrf=0.01, 47 | momentum=0.937, 48 | weight_decay=0.0005, 49 | warmup_epochs=3.0, 50 | warmup_momentum=0.8, 51 | warmup_bias_lr=0.1 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.015, 56 | hsv_s=0.7, 57 | hsv_v=0.4, 58 | degrees=0.0, 59 | translate=0.1, 60 | scale=0.9, 61 | shear=0.0, 62 | flipud=0.0, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.1, 66 | ) 67 | training_mode = "conv_silu" 68 | # use normal conv to speed up training and further improve accuracy. 69 | -------------------------------------------------------------------------------- /configs/yolov6l6.py: -------------------------------------------------------------------------------- 1 | # YOLOv6l6 model 2 | model = dict( 3 | type='YOLOv6l6', 4 | pretrained=None, 5 | depth_multiple=1.0, 6 | width_multiple=1.0, 7 | backbone=dict( 8 | type='CSPBepBackbone_P6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck_P6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512, 1024], 23 | num_layers=4, 24 | anchors=1, 25 | strides=[8, 16, 32, 64], 26 | atss_warmup_epoch=4, 27 | iou_type='giou', 28 | use_dfl=True, 29 | reg_max=16, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.01, 41 | lrf=0.01, 42 | momentum=0.937, 43 | weight_decay=0.0005, 44 | warmup_epochs=3.0, 45 | warmup_momentum=0.8, 46 | warmup_bias_lr=0.1 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.015, 51 | hsv_s=0.7, 52 | hsv_v=0.4, 53 | degrees=0.0, 54 | translate=0.1, 55 | scale=0.9, 56 | shear=0.0, 57 | flipud=0.0, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.2, 61 | ) 62 | training_mode = "conv_silu" -------------------------------------------------------------------------------- /configs/yolov6l6_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6l6 model 2 | model = dict( 3 | type='YOLOv6l6', 4 | pretrained='weights/yolov6l6.pt', 5 | depth_multiple=1.0, 6 | width_multiple=1.0, 7 | backbone=dict( 8 | type='CSPBepBackbone_P6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck_P6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512, 1024], 23 | num_layers=4, 24 | anchors=1, 25 | strides=[8, 16, 32, 64], 26 | atss_warmup_epoch=4, 27 | iou_type='giou', 28 | use_dfl=True, 29 | reg_max=16, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.0032, 41 | lrf=0.12, 42 | momentum=0.843, 43 | weight_decay=0.00036, 44 | warmup_epochs=2.0, 45 | warmup_momentum=0.5, 46 | warmup_bias_lr=0.05 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.0138, 51 | hsv_s=0.664, 52 | hsv_v=0.464, 53 | degrees=0.373, 54 | translate=0.245, 55 | scale=0.898, 56 | shear=0.602, 57 | flipud=0.00856, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.243, 61 | ) 62 | training_mode = "conv_silu" -------------------------------------------------------------------------------- /configs/yolov6l_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6l model 2 | model = dict( 3 | type='YOLOv6l', 4 | pretrained='weights/yolov6l.pt', 5 | depth_multiple=1.0, 6 | width_multiple=1.0, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(1)/2, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(1)/2, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 2.0, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver = dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.0032, 46 | lrf=0.12, 47 | momentum=0.843, 48 | weight_decay=0.00036, 49 | warmup_epochs=2.0, 50 | warmup_momentum=0.5, 51 | warmup_bias_lr=0.05 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.0138, 56 | hsv_s=0.664, 57 | hsv_v=0.464, 58 | degrees=0.373, 59 | translate=0.245, 60 | scale=0.898, 61 | shear=0.602, 62 | flipud=0.00856, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.243, 66 | ) 67 | training_mode = "conv_silu" 68 | # use normal conv to speed up training and further improve accuracy. 69 | -------------------------------------------------------------------------------- /configs/yolov6m.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m model 2 | model = dict( 3 | type='YOLOv6m', 4 | pretrained=None, 5 | depth_multiple=0.60, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(2)/3, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(2)/3, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 0.8, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver=dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.01, 46 | lrf=0.01, 47 | momentum=0.937, 48 | weight_decay=0.0005, 49 | warmup_epochs=3.0, 50 | warmup_momentum=0.8, 51 | warmup_bias_lr=0.1 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.015, 56 | hsv_s=0.7, 57 | hsv_v=0.4, 58 | degrees=0.0, 59 | translate=0.1, 60 | scale=0.9, 61 | shear=0.0, 62 | flipud=0.0, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.1, 66 | ) 67 | -------------------------------------------------------------------------------- /configs/yolov6m6.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m6 model 2 | model = dict( 3 | type='YOLOv6m6', 4 | pretrained=None, 5 | depth_multiple=0.60, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone_P6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | csp_e=float(2)/3, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck_P6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | csp_e=float(2)/3, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512, 1024], 23 | num_layers=4, 24 | anchors=1, 25 | strides=[8, 16, 32, 64], 26 | atss_warmup_epoch=4, 27 | iou_type='giou', 28 | use_dfl=True, 29 | reg_max=16, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.01, 41 | lrf=0.01, 42 | momentum=0.937, 43 | weight_decay=0.0005, 44 | warmup_epochs=3.0, 45 | warmup_momentum=0.8, 46 | warmup_bias_lr=0.1 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.015, 51 | hsv_s=0.7, 52 | hsv_v=0.4, 53 | degrees=0.0, 54 | translate=0.1, 55 | scale=0.9, 56 | shear=0.0, 57 | flipud=0.0, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.1, 61 | ) -------------------------------------------------------------------------------- /configs/yolov6m6_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m6 model 2 | model = dict( 3 | type='YOLOv6m6', 4 | pretrained='weights/yolov6m6.pt', 5 | depth_multiple=0.60, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone_P6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | csp_e=float(2)/3, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck_P6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | csp_e=float(2)/3, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512, 1024], 23 | num_layers=4, 24 | anchors=1, 25 | strides=[8, 16, 32, 64], 26 | atss_warmup_epoch=4, 27 | iou_type='giou', 28 | use_dfl=True, 29 | reg_max=16, #if use_dfl is False, please set reg_max to 0 30 | distill_weight={ 31 | 'class': 1.0, 32 | 'dfl': 1.0, 33 | }, 34 | ) 35 | ) 36 | 37 | solver = dict( 38 | optim='SGD', 39 | lr_scheduler='Cosine', 40 | lr0=0.0032, 41 | lrf=0.12, 42 | momentum=0.843, 43 | weight_decay=0.00036, 44 | warmup_epochs=2.0, 45 | warmup_momentum=0.5, 46 | warmup_bias_lr=0.05 47 | ) 48 | 49 | data_aug = dict( 50 | hsv_h=0.0138, 51 | hsv_s=0.664, 52 | hsv_v=0.464, 53 | degrees=0.373, 54 | translate=0.245, 55 | scale=0.898, 56 | shear=0.602, 57 | flipud=0.00856, 58 | fliplr=0.5, 59 | mosaic=1.0, 60 | mixup=0.243, 61 | ) -------------------------------------------------------------------------------- /configs/yolov6m_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6m model 2 | model = dict( 3 | type='YOLOv6m', 4 | pretrained='weights/yolov6m.pt', 5 | depth_multiple=0.60, 6 | width_multiple=0.75, 7 | backbone=dict( 8 | type='CSPBepBackbone', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | csp_e=float(2)/3, 12 | fuse_P2=True, 13 | ), 14 | neck=dict( 15 | type='CSPRepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | csp_e=float(2)/3, 19 | ), 20 | head=dict( 21 | type='EffiDeHead', 22 | in_channels=[128, 256, 512], 23 | num_layers=3, 24 | begin_indices=24, 25 | anchors=3, 26 | anchors_init=[[10,13, 19,19, 33,23], 27 | [30,61, 59,59, 59,119], 28 | [116,90, 185,185, 373,326]], 29 | out_indices=[17, 20, 23], 30 | strides=[8, 16, 32], 31 | atss_warmup_epoch=0, 32 | iou_type='giou', 33 | use_dfl=True, 34 | reg_max=16, #if use_dfl is False, please set reg_max to 0 35 | distill_weight={ 36 | 'class': 0.8, 37 | 'dfl': 1.0, 38 | }, 39 | ) 40 | ) 41 | 42 | solver = dict( 43 | optim='SGD', 44 | lr_scheduler='Cosine', 45 | lr0=0.0032, 46 | lrf=0.12, 47 | momentum=0.843, 48 | weight_decay=0.00036, 49 | warmup_epochs=2.0, 50 | warmup_momentum=0.5, 51 | warmup_bias_lr=0.05 52 | ) 53 | 54 | data_aug = dict( 55 | hsv_h=0.0138, 56 | hsv_s=0.664, 57 | hsv_v=0.464, 58 | degrees=0.373, 59 | translate=0.245, 60 | scale=0.898, 61 | shear=0.602, 62 | flipud=0.00856, 63 | fliplr=0.5, 64 | mosaic=1.0, 65 | mixup=0.243, 66 | ) 67 | -------------------------------------------------------------------------------- /configs/yolov6n.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='siou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.02, 45 | lrf=0.01, 46 | momentum=0.937, 47 | weight_decay=0.0005, 48 | warmup_epochs=3.0, 49 | warmup_momentum=0.8, 50 | warmup_bias_lr=0.1 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.015, 55 | hsv_s=0.7, 56 | hsv_v=0.4, 57 | degrees=0.0, 58 | translate=0.1, 59 | scale=0.5, 60 | shear=0.0, 61 | flipud=0.0, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.0, 65 | ) 66 | -------------------------------------------------------------------------------- /configs/yolov6n6.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n6', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | fuse_P2=True, # if use RepBiFPANNeck6, please set fuse_P2 to True. 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512, 1024], 22 | num_layers=4, 23 | anchors=1, 24 | strides=[8, 16, 32, 64], 25 | atss_warmup_epoch=4, 26 | iou_type='siou', 27 | use_dfl=False, 28 | reg_max=0 #if use_dfl is False, please set reg_max to 0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.02, 36 | lrf=0.01, 37 | momentum=0.937, 38 | weight_decay=0.0005, 39 | warmup_epochs=3.0, 40 | warmup_momentum=0.8, 41 | warmup_bias_lr=0.1 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.015, 46 | hsv_s=0.7, 47 | hsv_v=0.4, 48 | degrees=0.0, 49 | translate=0.1, 50 | scale=0.5, 51 | shear=0.0, 52 | flipud=0.0, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.0, 56 | ) 57 | -------------------------------------------------------------------------------- /configs/yolov6n6_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6n6', 4 | pretrained='weights/yolov6n6.pt', 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | fuse_P2=True, # if use RepBiFPANNeck6, please set fuse_P2 to True. 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512, 1024], 22 | num_layers=4, 23 | anchors=1, 24 | strides=[8, 16, 32, 64], 25 | atss_warmup_epoch=4, 26 | iou_type='siou', 27 | use_dfl=False, 28 | reg_max=0 #if use_dfl is False, please set reg_max to 0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.0032, 36 | lrf=0.12, 37 | momentum=0.843, 38 | weight_decay=0.00036, 39 | warmup_epochs=2.0, 40 | warmup_momentum=0.5, 41 | warmup_bias_lr=0.05 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.0138, 46 | hsv_s=0.664, 47 | hsv_v=0.464, 48 | degrees=0.373, 49 | translate=0.245, 50 | scale=0.898, 51 | shear=0.602, 52 | flipud=0.00856, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.243, 56 | ) 57 | -------------------------------------------------------------------------------- /configs/yolov6n_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6n', 4 | pretrained='weights/yolov6n.pt', 5 | depth_multiple=0.33, 6 | width_multiple=0.25, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='siou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.0032, 45 | lrf=0.12, 46 | momentum=0.843, 47 | weight_decay=0.00036, 48 | warmup_epochs=2.0, 49 | warmup_momentum=0.5, 50 | warmup_bias_lr=0.05 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.0138, 55 | hsv_s=0.664, 56 | hsv_v=0.464, 57 | degrees=0.373, 58 | translate=0.245, 59 | scale=0.898, 60 | shear=0.602, 61 | flipud=0.00856, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.243, 65 | ) 66 | -------------------------------------------------------------------------------- /configs/yolov6s.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6s', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='giou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.01, 45 | lrf=0.01, 46 | momentum=0.937, 47 | weight_decay=0.0005, 48 | warmup_epochs=3.0, 49 | warmup_momentum=0.8, 50 | warmup_bias_lr=0.1 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.015, 55 | hsv_s=0.7, 56 | hsv_v=0.4, 57 | degrees=0.0, 58 | translate=0.1, 59 | scale=0.5, 60 | shear=0.0, 61 | flipud=0.0, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.0, 65 | ) 66 | 67 | -------------------------------------------------------------------------------- /configs/yolov6s6.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6s6', 4 | pretrained=None, 5 | depth_multiple=0.33, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='EfficientRep6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | fuse_P2=True, # if use RepBiFPANNeck6, please set fuse_P2 to True. 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512, 1024], 22 | num_layers=4, 23 | anchors=1, 24 | strides=[8, 16, 32, 64], 25 | atss_warmup_epoch=4, 26 | iou_type='giou', 27 | use_dfl=False, 28 | reg_max=0 #if use_dfl is False, please set reg_max to 0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.01, 36 | lrf=0.01, 37 | momentum=0.937, 38 | weight_decay=0.0005, 39 | warmup_epochs=3.0, 40 | warmup_momentum=0.8, 41 | warmup_bias_lr=0.1 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.015, 46 | hsv_s=0.7, 47 | hsv_v=0.4, 48 | degrees=0.0, 49 | translate=0.1, 50 | scale=0.5, 51 | shear=0.0, 52 | flipud=0.0, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.0, 56 | ) 57 | -------------------------------------------------------------------------------- /configs/yolov6s6_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6n model 2 | model = dict( 3 | type='YOLOv6s6', 4 | pretrained='weights/yolov6s6.pt', 5 | depth_multiple=0.33, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='EfficientRep6', 9 | num_repeats=[1, 6, 12, 18, 6, 6], 10 | out_channels=[64, 128, 256, 512, 768, 1024], 11 | fuse_P2=True, # if use RepBiFPANNeck6, please set fuse_P2 to True. 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck6', 16 | num_repeats=[12, 12, 12, 12, 12, 12], 17 | out_channels=[512, 256, 128, 256, 512, 1024], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512, 1024], 22 | num_layers=4, 23 | anchors=1, 24 | strides=[8, 16, 32, 64], 25 | atss_warmup_epoch=4, 26 | iou_type='giou', 27 | use_dfl=False, 28 | reg_max=0 #if use_dfl is False, please set reg_max to 0 29 | ) 30 | ) 31 | 32 | solver = dict( 33 | optim='SGD', 34 | lr_scheduler='Cosine', 35 | lr0=0.0032, 36 | lrf=0.12, 37 | momentum=0.843, 38 | weight_decay=0.00036, 39 | warmup_epochs=2.0, 40 | warmup_momentum=0.5, 41 | warmup_bias_lr=0.05 42 | ) 43 | 44 | data_aug = dict( 45 | hsv_h=0.0138, 46 | hsv_s=0.664, 47 | hsv_v=0.464, 48 | degrees=0.373, 49 | translate=0.245, 50 | scale=0.898, 51 | shear=0.602, 52 | flipud=0.00856, 53 | fliplr=0.5, 54 | mosaic=1.0, 55 | mixup=0.243, 56 | ) -------------------------------------------------------------------------------- /configs/yolov6s_finetune.py: -------------------------------------------------------------------------------- 1 | # YOLOv6s model 2 | model = dict( 3 | type='YOLOv6s', 4 | pretrained='weights/yolov6s.pt', 5 | depth_multiple=0.33, 6 | width_multiple=0.50, 7 | backbone=dict( 8 | type='EfficientRep', 9 | num_repeats=[1, 6, 12, 18, 6], 10 | out_channels=[64, 128, 256, 512, 1024], 11 | fuse_P2=True, 12 | cspsppf=True, 13 | ), 14 | neck=dict( 15 | type='RepBiFPANNeck', 16 | num_repeats=[12, 12, 12, 12], 17 | out_channels=[256, 128, 128, 256, 256, 512], 18 | ), 19 | head=dict( 20 | type='EffiDeHead', 21 | in_channels=[128, 256, 512], 22 | num_layers=3, 23 | begin_indices=24, 24 | anchors=3, 25 | anchors_init=[[10,13, 19,19, 33,23], 26 | [30,61, 59,59, 59,119], 27 | [116,90, 185,185, 373,326]], 28 | out_indices=[17, 20, 23], 29 | strides=[8, 16, 32], 30 | atss_warmup_epoch=0, 31 | iou_type='giou', 32 | use_dfl=False, # set to True if you want to further train with distillation 33 | reg_max=0, # set to 16 if you want to further train with distillation 34 | distill_weight={ 35 | 'class': 1.0, 36 | 'dfl': 1.0, 37 | }, 38 | ) 39 | ) 40 | 41 | solver = dict( 42 | optim='SGD', 43 | lr_scheduler='Cosine', 44 | lr0=0.0032, 45 | lrf=0.12, 46 | momentum=0.843, 47 | weight_decay=0.00036, 48 | warmup_epochs=2.0, 49 | warmup_momentum=0.5, 50 | warmup_bias_lr=0.05 51 | ) 52 | 53 | data_aug = dict( 54 | hsv_h=0.0138, 55 | hsv_s=0.664, 56 | hsv_v=0.464, 57 | degrees=0.373, 58 | translate=0.245, 59 | scale=0.898, 60 | shear=0.602, 61 | flipud=0.00856, 62 | fliplr=0.5, 63 | mosaic=1.0, 64 | mixup=0.243, 65 | ) -------------------------------------------------------------------------------- /data/coco.yaml: -------------------------------------------------------------------------------- 1 | # COCO 2017 dataset http://cocodataset.org 2 | train: ../coco/images/train2017 # 118287 images 3 | val: ../coco/images/val2017 # 5000 images 4 | test: ../coco/images/test2017 5 | anno_path: ../coco/annotations/instances_val2017.json 6 | 7 | # number of classes 8 | nc: 80 9 | # whether it is coco dataset, only coco dataset should be set to True. 10 | is_coco: True 11 | 12 | # class names 13 | names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 14 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 15 | 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 16 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 17 | 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 18 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 19 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 20 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 21 | 'hair drier', 'toothbrush' ] 22 | -------------------------------------------------------------------------------- /data/dataset.yaml: -------------------------------------------------------------------------------- 1 | # Please insure that your custom_dataset are put in same parent dir with YOLOv6_DIR 2 | train: ../custom_dataset/images/train # train images 3 | val: ../custom_dataset/images/val # val images 4 | test: ../custom_dataset/images/test # test images (optional) 5 | 6 | # whether it is coco dataset, only coco dataset should be set to True. 7 | is_coco: False 8 | # Classes 9 | nc: 20 # number of classes 10 | names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 11 | 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'] # class names 12 | -------------------------------------------------------------------------------- /data/images/image1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/data/images/image1.jpg -------------------------------------------------------------------------------- /data/images/image2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/data/images/image2.jpg -------------------------------------------------------------------------------- /data/images/image3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/data/images/image3.jpg -------------------------------------------------------------------------------- /data/voc.yaml: -------------------------------------------------------------------------------- 1 | # Please insure that your custom_dataset are put in same parent dir with YOLOv6_DIR 2 | train: VOCdevkit/voc_07_12/images/train # train images 3 | val: VOCdevkit/voc_07_12/images/val # val images 4 | test: VOCdevkit/voc_07_12/images/val # test images (optional) 5 | 6 | # whether it is coco dataset, only coco dataset should be set to True. 7 | is_coco: False 8 | # Classes 9 | nc: 20 # number of classes 10 | names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 11 | 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'] # class names 12 | -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/README.md: -------------------------------------------------------------------------------- 1 | # Object Detection using YOLOv5/YOLOv6/YOLOX and OpenCV DNN (Python/C++) 2 | 3 | ## 0. Install Dependancies 4 | ``` 5 | OpenCV >= 4.5.4 6 | ``` 7 | Only **OpenCV >= 4.5.4** can read onnx model file by dnn module. 8 | 9 | ## 1. Usage 10 | Change work directory to `/path/to/YOLOv6/deploy/ONNX/OpenCV` 11 | ### 1.1 Python 12 | 13 | - YOLOv5 & YOLOv6: 14 | ```Python 15 | python yolo.py --model /path/to/onnx/yolov5n.onnx --img /path/to/sample.jpg --classesFile /path/to/coco.names 16 | yolov5s.onnx 17 | yolov5m.onnx 18 | yolov6n.onnx 19 | yolov6s.onnx 20 | yolov6t.onnx 21 | ``` 22 | - YOLOX: 23 | ```Python 24 | python yolox.py --model /path/to/onnx/yolox_nano.onnx --img /path/to/sample.jpg --classesFile /path/to/coco.names 25 | yolox_tiny.onnx 26 | yolox_s.onnx 27 | yolox_m.onnx 28 | ``` 29 | 30 | ### 1.2 CMake C++ Linux YOLOv5 31 | ```C++ Linux 32 | cd yolov5 // modify CMakeLists.txt 33 | mkdir build 34 | cd build 35 | cmake .. 36 | make 37 | ./yolov5 /path/to/onnx/yolov5n.onnx /path/to/sample.jpg /path/to/coco.names 38 | yolov5s.onnx 39 | yolov5m.onnx 40 | ``` 41 | 42 | ### 1.3 CMake C++ Linux YOLOv6 43 | ```C++ Linux 44 | cd yolov6 // modify CMakeLists.txt 45 | mkdir build 46 | cd build 47 | cmake .. 48 | make 49 | ./yolov6 /path/to/onnx/yolov6n.onnx /path/to/sample.jpg /path/to/coco.names 50 | yolov6t.onnx 51 | yolov6s.onnx 52 | ``` 53 | 54 | ### 1.4 CMake C++ Linux YOLOX 55 | ```C++ Linux 56 | cd yolox // modify CMakeLists.txt 57 | mkdir build 58 | cd build 59 | cmake .. 60 | make 61 | ./yolox /path/to/onnx/yolox_nano.onnx /path/to/sample.jpg /path/to/coco.names 62 | yolox_tiny.onnx 63 | yolox_s.onnx 64 | yolox_m.onnx 65 | ``` 66 | 67 | ## 2. Result 68 | | Model | Speed CPU b1(ms) Python | Speed CPU b1(ms) C++ | mAPval 0.5:0.95 | params(M) | FLOPs(G) | 69 | | :-- | :-: | :-: | :-: | :-: | :-: | 70 | | **YOLOv5n** | 116.47 | 118.89 | 28.0 | 1.9 | 4.5 | 71 | | **YOLOv5s** | 200.53 | 202.22 | 37.4 | 7.2 | 16.5 | 72 | | **YOLOv5m** | 294.98 | 291.86 | 45.4 | 21.2 | 49.0 | 73 | | | | | | | | 74 | | **YOLOv6-n** | 62.37 | 60.34 | 37.5 | 4.7 | 11.4 | 75 | | **YOLOv6-s** | 137.94 | 148.01 | 45.0 | 18.5 | 45.3 | 76 | | **YOLOv6-m** | 264.40 | 269.31 | 50.0 | 34.9 | 85.8 | 77 | | | | | | | | 78 | | **YOLOX-Nano** | 81.06 | 86.75 | 25.8@416 | 0.91 | 1.08@416 | 79 | | **YOLOX-tiny** | 129.72 | 144.19 | 32.8@416 | 5.06 | 6.45@416 | 80 | | **YOLOX-s** | 180.86 | 169.96 | 40.5 | 9.0 | 26.8 | 81 | | **YOLOX-m** | 336.34 | 357.91 | 47.2 | 25.3 | 73.8 | 82 | 83 | **Note**: 84 | - All onnx models are converted from official github([Google Drive](https://drive.google.com/drive/folders/1Nw6M_Y6XLASyB0RxhSI2z_QRtt70Picl?usp=sharing)). 85 | - Speed is test by [dnn::Net::getPerfProfile](https://docs.opencv.org/4.5.5/db/d30/classcv_1_1dnn_1_1Net.html), we report the average inference time of 300 runs on the same environment. 86 | - The mAP/params/FLOPs are from official github. 87 | - Test environment: MacOS 11.4 with 2.6 GHz 6-core Intel Core i7, 16GB Memory. 88 | 89 | ### Visualization 90 |
91 |
92 |
93 | -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/coco.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | aeroplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | sofa 59 | pottedplant 60 | bed 61 | diningtable 62 | toilet 63 | tvmonitor 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush 81 | -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/sample.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/deploy/ONNX/OpenCV/sample.jpg -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/yolo.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import os 4 | import argparse 5 | 6 | 7 | # Constants. 8 | INPUT_WIDTH = 640 9 | INPUT_HEIGHT = 640 10 | SCORE_THRESHOLD = 0.5 # cls score 11 | NMS_THRESHOLD = 0.45 12 | CONFIDENCE_THRESHOLD = 0.45 # obj confidence 13 | 14 | # Text parameters. 15 | FONT_FACE = cv2.FONT_HERSHEY_SIMPLEX 16 | FONT_SCALE = 0.7 17 | THICKNESS = 1 18 | 19 | # Colors 20 | BLACK = (0,0,0) 21 | BLUE = (255,178,50) 22 | YELLOW = (0,255,255) 23 | RED = (0,0,255) 24 | 25 | 26 | def draw_label(input_image, label, left, top): 27 | """Draw text onto image at location.""" 28 | 29 | # Get text size. 30 | text_size = cv2.getTextSize(label, FONT_FACE, FONT_SCALE, THICKNESS) 31 | dim, baseline = text_size[0], text_size[1] 32 | # Use text size to create a BLACK rectangle. 33 | cv2.rectangle(input_image, (left, top), (left + dim[0], top + dim[1] + baseline), BLACK, cv2.FILLED) 34 | # Display text inside the rectangle. 35 | cv2.putText(input_image, label, (left, top + dim[1]), FONT_FACE, FONT_SCALE, YELLOW, THICKNESS, cv2.LINE_AA) 36 | 37 | 38 | def pre_process(input_image, net): 39 | # Create a 4D blob from a frame. 40 | blob = cv2.dnn.blobFromImage(input_image, 1/255, (INPUT_WIDTH, INPUT_HEIGHT), [0,0,0], 1, crop=False) 41 | 42 | # Sets the input to the network. 43 | net.setInput(blob) 44 | 45 | # Runs the forward pass to get output of the output layers. 46 | output_layers = net.getUnconnectedOutLayersNames() 47 | outputs = net.forward(output_layers) 48 | # print(outputs[0].shape) 49 | 50 | return outputs 51 | 52 | 53 | def post_process(input_image, outputs): 54 | # Lists to hold respective values while unwrapping. 55 | class_ids = [] 56 | confidences = [] 57 | boxes = [] 58 | 59 | # Rows. 60 | rows = outputs[0].shape[1] 61 | 62 | image_height, image_width = input_image.shape[:2] 63 | 64 | # Resizing factor. 65 | x_factor = image_width / INPUT_WIDTH 66 | y_factor = image_height / INPUT_HEIGHT 67 | 68 | # Iterate through 25200 detections. 69 | for r in range(rows): 70 | row = outputs[0][0][r] 71 | confidence = row[4] 72 | 73 | # Discard bad detections and continue. 74 | if confidence >= CONFIDENCE_THRESHOLD: 75 | classes_scores = row[5:] 76 | 77 | # Get the index of max class score. 78 | class_id = np.argmax(classes_scores) 79 | 80 | # Continue if the class score is above threshold. 81 | if (classes_scores[class_id] > SCORE_THRESHOLD): 82 | confidences.append(confidence) 83 | class_ids.append(class_id) 84 | 85 | cx, cy, w, h = row[0], row[1], row[2], row[3] 86 | 87 | left = int((cx - w/2) * x_factor) 88 | top = int((cy - h/2) * y_factor) 89 | width = int(w * x_factor) 90 | height = int(h * y_factor) 91 | 92 | box = np.array([left, top, width, height]) 93 | boxes.append(box) 94 | 95 | # Perform non maximum suppression to eliminate redundant overlapping boxes with 96 | # lower confidences. 97 | indices = cv2.dnn.NMSBoxes(boxes, confidences, CONFIDENCE_THRESHOLD, NMS_THRESHOLD) 98 | for i in indices: 99 | box = boxes[i] 100 | left = box[0] 101 | top = box[1] 102 | width = box[2] 103 | height = box[3] 104 | cv2.rectangle(input_image, (left, top), (left + width, top + height), BLUE, 3*THICKNESS) 105 | label = "{}:{:.2f}".format(classes[class_ids[i]], confidences[i]) 106 | draw_label(input_image, label, left, top) 107 | 108 | return input_image 109 | 110 | 111 | if __name__ == '__main__': 112 | parser = argparse.ArgumentParser() 113 | parser.add_argument('--model', default='models/yolov6n.onnx', help="Input your onnx model.") 114 | parser.add_argument('--img', default='sample.jpg', help="Path to your input image.") 115 | parser.add_argument('--classesFile', default='coco.names', help="Path to your classesFile.") 116 | args = parser.parse_args() 117 | 118 | # Load class names. 119 | model_path, img_path, classesFile = args.model, args.img, args.classesFile 120 | window_name = os.path.splitext(os.path.basename(model_path))[0] 121 | classes = None 122 | with open(classesFile, 'rt') as f: 123 | classes = f.read().rstrip('\n').split('\n') 124 | 125 | # Load image. 126 | frame = cv2.imread(img_path) 127 | input = frame.copy() 128 | 129 | # Give the weight files to the model and load the network using them. 130 | net = cv2.dnn.readNet(model_path) 131 | 132 | # Put efficiency information. The function getPerfProfile returns the overall time for inference(t) and the 133 | # timings for each of the layers(in layersTimes) 134 | # Process image. 135 | cycles = 300 136 | total_time = 0 137 | for i in range(cycles): 138 | detections = pre_process(input.copy(), net) 139 | img = post_process(frame.copy(), detections) 140 | t, _ = net.getPerfProfile() 141 | total_time += t 142 | print(f'Cycle [{i + 1}]:\t{t * 1000.0 / cv2.getTickFrequency():.2f}\tms') 143 | 144 | avg_time = total_time / cycles 145 | label = 'Average Inference time: %.2f ms' % (avg_time * 1000.0 / cv2.getTickFrequency()) 146 | print(f'Model: {window_name}\n{label}') 147 | cv2.putText(img, label, (20, 40), FONT_FACE, FONT_SCALE, RED, THICKNESS, cv2.LINE_AA) 148 | cv2.imshow(window_name, img) 149 | cv2.waitKey(0) 150 | -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/yolov5/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # CMakeLists.txt 2 | 3 | # Older versions of CMake are likely to work just fine but, since 4 | # I don't know where to cut off I just use the version I'm using 5 | cmake_minimum_required(VERSION "3.17") 6 | 7 | # name of this example project 8 | project(simple-demo) 9 | 10 | # set OpenCV_DIR variable equal to the path to the cmake 11 | # files within the previously installed opencv program 12 | # path like /xxx/yyy/opencv/install/lib/cmake/opencv4 13 | set(OpenCV_DIR ${OpenCV_DIR}) 14 | 15 | # Tell compiler to use C++ 14 features which is needed because 16 | # Clang version is often behind in the XCode installation 17 | set(CMAKE_CXX_STANDARD 14) 18 | 19 | # configure the necessary common CMake environment variables 20 | # needed to include and link the OpenCV program into this 21 | # demo project, namely OpenCV_INCLUDE_DIRS and OpenCV_LIBS 22 | find_package( OpenCV REQUIRED ) 23 | 24 | # tell the build to include the headers from OpenCV 25 | include_directories( ${OpenCV_INCLUDE_DIRS} ) 26 | 27 | # specify the executable target to be built 28 | # path like /xxx/yyy/opencv/install/include 29 | add_executable(yolov5 yolov5.cpp) 30 | 31 | # tell it to link the executable target against OpenCV 32 | # path like /xxx/yyy/opencv/install/lib 33 | target_link_libraries(yolov5 ${OpenCV_LIBS} ) 34 | -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/yolov6/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # CMakeLists.txt 2 | 3 | # Older versions of CMake are likely to work just fine but, since 4 | # I don't know where to cut off I just use the version I'm using 5 | cmake_minimum_required(VERSION "3.17") 6 | 7 | # name of this example project 8 | project(simple-demo) 9 | 10 | # set OpenCV_DIR variable equal to the path to the cmake 11 | # files within the previously installed opencv program 12 | # path like /xxx/yyy/opencv/install/lib/cmake/opencv4 13 | set(OpenCV_DIR ${OpenCV_DIR}) 14 | 15 | # Tell compiler to use C++ 14 features which is needed because 16 | # Clang version is often behind in the XCode installation 17 | set(CMAKE_CXX_STANDARD 14) 18 | 19 | # configure the necessary common CMake environment variables 20 | # needed to include and link the OpenCV program into this 21 | # demo project, namely OpenCV_INCLUDE_DIRS and OpenCV_LIBS 22 | find_package( OpenCV REQUIRED ) 23 | 24 | # tell the build to include the headers from OpenCV 25 | # path like /xxx/yyy/opencv/install/include 26 | include_directories( ${OpenCV_INCLUDE_DIRS} ) 27 | 28 | # specify the executable target to be built 29 | add_executable(yolov6 yolov6.cpp) 30 | 31 | # tell it to link the executable target against OpenCV 32 | # path like /xxx/yyy/opencv/install/lib 33 | target_link_libraries(yolov6 ${OpenCV_LIBS} ) 34 | -------------------------------------------------------------------------------- /deploy/ONNX/OpenCV/yolox/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # CMakeLists.txt 2 | 3 | # Older versions of CMake are likely to work just fine but, since 4 | # I don't know where to cut off I just use the version I'm using 5 | cmake_minimum_required(VERSION "3.17") 6 | 7 | # name of this example project 8 | project(simple-demo) 9 | 10 | # set OpenCV_DIR variable equal to the path to the cmake 11 | # files within the previously installed opencv program 12 | # path like /xxx/yyy/opencv/install/lib/cmake/opencv4 13 | set(OpenCV_DIR ${OpenCV_DIR}) 14 | 15 | # Tell compiler to use C++ 14 features which is needed because 16 | # Clang version is often behind in the XCode installation 17 | set(CMAKE_CXX_STANDARD 14) 18 | 19 | # configure the necessary common CMake environment variables 20 | # needed to include and link the OpenCV program into this 21 | # demo project, namely OpenCV_INCLUDE_DIRS and OpenCV_LIBS 22 | find_package( OpenCV REQUIRED ) 23 | 24 | # tell the build to include the headers from OpenCV 25 | # path like /xxx/yyy/opencv/install/include 26 | include_directories( ${OpenCV_INCLUDE_DIRS} ) 27 | 28 | # specify the executable target to be built 29 | add_executable(yolox yolox.cpp) 30 | 31 | # tell it to link the executable target against OpenCV 32 | # path like /xxx/yyy/opencv/install/lib 33 | target_link_libraries(yolox ${OpenCV_LIBS} ) 34 | -------------------------------------------------------------------------------- /deploy/ONNX/README.md: -------------------------------------------------------------------------------- 1 | # Export ONNX Model 2 | 3 | ## Check requirements 4 | ```shell 5 | pip install onnx>=1.10.0 6 | ``` 7 | 8 | ## Export script 9 | ```shell 10 | python ./deploy/ONNX/export_onnx.py \ 11 | --weights yolov6s.pt \ 12 | --img 640 \ 13 | --batch 1 \ 14 | --simplify 15 | ``` 16 | 17 | 18 | 19 | #### Description of all arguments 20 | 21 | - `--weights` : The path of yolov6 model weights. 22 | - `--img` : Image size of model inputs. 23 | - `--batch` : Batch size of model inputs. 24 | - `--half` : Whether to export half-precision model. 25 | - `--inplace` : Whether to set Detect() inplace. 26 | - `--simplify` : Whether to simplify onnx. Not support in end to end export. 27 | - `--end2end` : Whether to export end to end onnx model. Only support onnxruntime and TensorRT >= 8.0.0 . 28 | - `--trt-version` : Export onnx for TensorRT version. Support : 7 or 8. 29 | - `--ort` : Whether to export onnx for onnxruntime backend. 30 | - `--with-preprocess` : Whether to export preprocess with bgr2rgb and normalize (divide by 255) 31 | - `--topk-all` : Topk objects for every image. 32 | - `--iou-thres` : IoU threshold for NMS algorithm. 33 | - `--conf-thres` : Confidence threshold for NMS algorithm. 34 | - `--device` : Export device. Cuda device : 0 or 0,1,2,3 ... , CPU : cpu . 35 | 36 | ## Download 37 | 38 | * [YOLOv6-N](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6n.onnx) 39 | * [YOLOv6-T](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6t.onnx) 40 | * [YOLOv6-S](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6s.onnx) 41 | * [YOLOv6-M](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6m.onnx) 42 | * [YOLOv6-L-ReLU](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6l_relu.onnx) 43 | * [YOLOv6-L](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6l.onnx) 44 | 45 | 46 | ## End2End export 47 | 48 | Now YOLOv6 supports end to end detect for onnxruntime and TensorRT ! 49 | 50 | If you want to deploy in TensorRT, make sure you have installed TensorRT ! 51 | 52 | ### onnxruntime backend 53 | #### Usage 54 | 55 | ```bash 56 | python ./deploy/ONNX/export_onnx.py \ 57 | --weights yolov6s.pt \ 58 | --img 640 \ 59 | --batch 1 \ 60 | --end2end \ 61 | --ort 62 | ``` 63 | 64 | You will get an onnx with **NonMaxSuppression** operator . 65 | 66 | ### TensorRT backend (TensorRT version == 7.2.3.4) 67 | #### Usage 68 | ```bash 69 | python ./deploy/ONNX/export_onnx.py \ 70 | --weights yolov6s.pt \ 71 | --img 640 \ 72 | --batch 1 \ 73 | --end2end \ 74 | --trt-version 7 75 | ``` 76 | You will get an onnx with **[BatchedNMSDynamic_TRT](https://github.com/triple-Mu/TensorRT/tree/main/plugin/batchedNMSPlugin)** plugin . 77 | 78 | 79 | ### TensorRT backend (TensorRT version>= 8.0.0) 80 | 81 | #### Usage 82 | 83 | ```bash 84 | python ./deploy/ONNX/export_onnx.py \ 85 | --weights yolov6s.pt \ 86 | --img 640 \ 87 | --batch 1 \ 88 | --end2end \ 89 | --trt-version 8 90 | ``` 91 | 92 | You will get an onnx with **[EfficientNMS_TRT](https://github.com/NVIDIA/TensorRT/tree/main/plugin/efficientNMSPlugin)** plugin . 93 | 94 | ### Outputs Description 95 | 96 | The onnx outputs are as shown : 97 | 98 | 99 | 100 | ```num_dets``` means the number of object in every image in its batch . 101 | 102 | ```det_boxes``` means topk(100) object's location about [`x0`,`y0`,`x1`,`y1`] . 103 | 104 | ```det_scores``` means the confidence score of every topk(100) objects . 105 | 106 | ```det_classes``` means the category of every topk(100) objects . 107 | 108 | 109 | You can export TensorRT engine use [trtexec](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#trtexec-ovr) tools. 110 | #### Usage 111 | For both TensorRT-7 and TensorRT-8 `trtexec` tool is avaiable. 112 | ``` shell 113 | trtexec --onnx=yolov6s.onnx \ 114 | --saveEngine=yolov6s.engine \ 115 | --workspace=8192 # 8GB 116 | --fp16 # if export TensorRT fp16 model 117 | ``` 118 | 119 | ## Evaluate TensorRT model's performance 120 | 121 | When we get the TensorRT model, we can evalute its performance by: 122 | ``` 123 | python deploy/ONNX/eval_trt.py --weights yolov6s.engine --batch-size=1 --data data/coco.yaml 124 | ``` 125 | 126 | ## Dynamic Batch Inference 127 | 128 | YOLOv6 support dynamic batch export and inference, you can refer to: 129 | 130 | [export ONNX model with dynamic batch ](YOLOv6-Dynamic-Batch-onnxruntime.ipynb) 131 | 132 | [export TensorRT model with dynamic batch](YOLOv6-Dynamic-Batch-tensorrt.ipynb) 133 | -------------------------------------------------------------------------------- /deploy/ONNX/eval_trt.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import argparse 4 | import os 5 | import os.path as osp 6 | import sys 7 | import torch 8 | 9 | ROOT = os.getcwd() 10 | if str(ROOT) not in sys.path: 11 | sys.path.append(str(ROOT)) 12 | 13 | from yolov6.core.evaler import Evaler 14 | from yolov6.utils.events import LOGGER 15 | from yolov6.utils.general import increment_name 16 | 17 | 18 | def get_args_parser(add_help=True): 19 | parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Evalating', add_help=add_help) 20 | parser.add_argument('--data', type=str, default='./data/coco.yaml', help='dataset yaml file path.') 21 | parser.add_argument('--weights', type=str, default='./yolov6s.engine', help='tensorrt engine file path.') 22 | parser.add_argument('--batch-size', type=int, default=32, help='batch size') 23 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 24 | parser.add_argument('--task', default='val', help='can only be val now.') 25 | parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 26 | parser.add_argument('--save_dir', type=str, default='runs/val/', help='evaluation save dir') 27 | parser.add_argument('--name', type=str, default='exp', help='save evaluation results to save_dir/name') 28 | args = parser.parse_args() 29 | LOGGER.info(args) 30 | return args 31 | 32 | 33 | @torch.no_grad() 34 | def run(data, 35 | weights=None, 36 | batch_size=32, 37 | img_size=640, 38 | task='val', 39 | device='', 40 | save_dir='', 41 | name = '' 42 | ): 43 | """ 44 | TensorRT models's evaluation process. 45 | """ 46 | 47 | # task 48 | assert task== 'val', f'task type can only be val, however you set it to {task}' 49 | 50 | save_dir = str(increment_name(osp.join(save_dir, name))) 51 | os.makedirs(save_dir, exist_ok=True) 52 | 53 | dummy_model = torch.zeros(0) 54 | device = Evaler.reload_device(device, dummy_model, task) 55 | 56 | data = Evaler.reload_dataset(data) if isinstance(data, str) else data 57 | 58 | # init 59 | val = Evaler(data, batch_size, img_size, None, \ 60 | None, device, False, save_dir) 61 | 62 | dataloader,pred_result = val.eval_trt(weights) 63 | eval_result = val.eval_model(pred_result, dummy_model, dataloader, task) 64 | return eval_result 65 | 66 | 67 | def main(args): 68 | run(**vars(args)) 69 | 70 | 71 | if __name__ == "__main__": 72 | args = get_args_parser() 73 | main(args) 74 | -------------------------------------------------------------------------------- /deploy/OpenVINO/README.md: -------------------------------------------------------------------------------- 1 | ## Export OpenVINO Model 2 | 3 | ### Check requirements 4 | ```shell 5 | pip install --upgrade pip 6 | pip install openvino-dev 7 | ``` 8 | 9 | ### Export script 10 | ```shell 11 | python deploy/OpenVINO/export_openvino.py --weights yolov6s.pt --img 640 --batch 1 12 | 13 | ``` 14 | 15 | ### Speed test 16 | ```shell 17 | benchmark_app -m yolov6s_openvino/yolov6s.xml -i data/images/image1.jpg -d CPU -niter 100 -progress 18 | 19 | ``` 20 | -------------------------------------------------------------------------------- /deploy/OpenVINO/export_openvino.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import argparse 4 | import time 5 | import sys 6 | import os 7 | import torch 8 | import torch.nn as nn 9 | import onnx 10 | import subprocess 11 | 12 | ROOT = os.getcwd() 13 | if str(ROOT) not in sys.path: 14 | sys.path.append(str(ROOT)) 15 | 16 | from yolov6.models.yolo import * 17 | from yolov6.models.effidehead import Detect 18 | from yolov6.layers.common import * 19 | from yolov6.utils.events import LOGGER 20 | from yolov6.utils.checkpoint import load_checkpoint 21 | 22 | 23 | if __name__ == '__main__': 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--weights', type=str, default='./yolov6s.pt', help='weights path') 26 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width 27 | parser.add_argument('--batch-size', type=int, default=1, help='batch size') 28 | parser.add_argument('--half', action='store_true', help='FP16 half-precision export') 29 | parser.add_argument('--inplace', action='store_true', help='set Detect() inplace=True') 30 | parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 31 | args = parser.parse_args() 32 | args.img_size *= 2 if len(args.img_size) == 1 else 1 # expand 33 | print(args) 34 | t = time.time() 35 | 36 | # Check device 37 | cuda = args.device != 'cpu' and torch.cuda.is_available() 38 | device = torch.device('cuda:0' if cuda else 'cpu') 39 | assert not (device.type == 'cpu' and args.half), '--half only compatible with GPU export, i.e. use --device 0' 40 | # Load PyTorch model 41 | model = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True) # load FP32 model 42 | for layer in model.modules(): 43 | if isinstance(layer, RepVGGBlock): 44 | layer.switch_to_deploy() 45 | 46 | # Input 47 | img = torch.zeros(args.batch_size, 3, *args.img_size).to(device) # image size(1,3,320,192) iDetection 48 | 49 | # Update model 50 | if args.half: 51 | img, model = img.half(), model.half() # to FP16 52 | model.eval() 53 | for k, m in model.named_modules(): 54 | if isinstance(m, Conv): # assign export-friendly activations 55 | if isinstance(m.act, nn.SiLU): 56 | m.act = SiLU() 57 | elif isinstance(m, Detect): 58 | m.inplace = args.inplace 59 | 60 | y = model(img) # dry run 61 | 62 | # ONNX export 63 | try: 64 | LOGGER.info('\nStarting to export ONNX...') 65 | export_file = args.weights.replace('.pt', '.onnx') # filename 66 | torch.onnx.export(model, img, export_file, verbose=False, opset_version=12, 67 | training=torch.onnx.TrainingMode.EVAL, 68 | do_constant_folding=True, 69 | input_names=['image_arrays'], 70 | output_names=['outputs'], 71 | ) 72 | 73 | # Checks 74 | onnx_model = onnx.load(export_file) # load onnx model 75 | onnx.checker.check_model(onnx_model) # check onnx model 76 | LOGGER.info(f'ONNX export success, saved as {export_file}') 77 | except Exception as e: 78 | LOGGER.info(f'ONNX export failure: {e}') 79 | 80 | # OpenVINO export 81 | try: 82 | LOGGER.info('\nStarting to export OpenVINO...') 83 | import_file = args.weights.replace('.pt', '.onnx') 84 | export_dir = str(import_file).replace('.onnx', '_openvino') 85 | cmd = f"mo --input_model {import_file} --output_dir {export_dir} --data_type {'FP16' if args.half else 'FP32'}" 86 | subprocess.check_output(cmd.split()) 87 | LOGGER.info(f'OpenVINO export success, saved as {export_dir}') 88 | except Exception as e: 89 | LOGGER.info(f'OpenVINO export failure: {e}') 90 | 91 | # Finish 92 | LOGGER.info('\nExport complete (%.2fs)' % (time.time() - t)) 93 | -------------------------------------------------------------------------------- /deploy/RKNN/RKOPT_README.md: -------------------------------------------------------------------------------- 1 | ## Description - export optimized model for RKNPU 2 | 3 | ### 1. Model structure Adjustment 4 | 5 | - The dfl structure has poor performance on NPU processing, moved outside the model. 6 | 7 | Assuming that there are 6000 candidate frames, the original model places the dfl structure before the "box confidence filter", then the 6000 candidate frames need to be calculated through dfl calculation. If the dfl structure is placed after the "box confidence filter", Assuming that there are 100 candidate boxes left after filtering, the calculation amount of the dfl part is reduced to 100, which greatly reduces the occupancy of computing resources and bandwidth resources. 8 | 9 | - Notice: yolov6n/s hasn't dfl structure, while yolov6m/l has dfl structure 10 | 11 | 12 | 13 | - Assuming that there are 6000 candidate boxes and the detection category is 80, the threshold retrieval operation needs to be repeated 6000* 80 ~= 4.8*10^5 times, which takes a lot of time. Therefore, when exporting the model, an additional summation operation for 80 types of detection targets is added to the model to quickly filter the confidence. (This structure is effective in some cases, related to the training results of the model) 14 | 15 | (v6m, v6l) To disable this optimization, comment the following code in ./yolov6/models/effidehead.py (line70~86 part) 16 | 17 | ``` 18 | cls_sum = torch.clamp(y[-1].sum(1, keepdim=True), 0, 1) 19 | output_for_rknn.append(cls_sum) 20 | ``` 21 | 22 | (v6n, v6s) To disable this optimization, comment the following code in ./yolov6/models/heads/effidehead_distill_ns.py (line78~94 part) 23 | 24 | ``` 25 | cls_sum = torch.clamp(y[-1].sum(1, keepdim=True), 0, 1) 26 | output_for_rknn.append(cls_sum) 27 | ``` 28 | 29 | 30 | 31 | 32 | - This optimization only affects the export of the model and does not affect the training process. **For the training steps, please refer to the YOLOv6 official documentation.** 33 | 34 | 35 | 36 | ### 2. Export model operation 37 | 38 | After meeting the environmental requirements of ./requirements.txt, execute the following statement to export the model 39 | 40 | ``` 41 | python deploy/RKNN/export_onnx_for_rknn.py --weight ./yolov6n.pt 42 | 43 | # adjust ./yolov6n.pt path to export your model. 44 | ``` 45 | 46 | 47 | 48 | ### 3.Transfer to RKNN model, Python demo, C demo 49 | 50 | Please refer https://github.com/airockchip/rknn_model_zoo/tree/main/models/CV/object_detection/yolo -------------------------------------------------------------------------------- /deploy/RKNN/RKOPT_README_cn.md: -------------------------------------------------------------------------------- 1 | ## RKNN 导出模型说明 2 | 3 | ### 1.调整部分 4 | 5 | - 由于 dfl 结构在 npu 处理性能不佳。假设有6000个候选框，原模型将 dfl 结构放置于 ''框置信度过滤" 前，则 6000 个候选框都需要计算经过 dfl 计算；而将 dfl 结构放置于 ''框置信度过滤" 后，假设过程成 100 个候选框，则dfl部分计算量减少至 100 个。 6 | 7 | 故将 dfl 结构使用 cpu 处理的耗时，虽然享受不到 npu 加速，但是本来带来的计算量较少也是很可观的。 8 | 9 | 注： yolov6n, yolov6s 没有 dfl 结构; yolov6m, yolov6l 存在 dfl 结构 10 | 11 | 12 | 13 | - 假设存在 6000 个候选框，存在 80 类检测目标，则阈值需要检索的置信度有 6000* 80 ～= 4.8*10^5 个，占据了较多耗时，故导出模型时，在模型中额外新增了对 80 类检测目标进行求和操作，用于快速过滤置信度，该结构在部分情况下对模型有效。 14 | 15 | (v6m, v6l) 可以在 ./yolov6/models/effidehead.py 70~86行位置，注释掉这部分 16 | 17 | ``` 18 | cls_sum = torch.clamp(y[-1].sum(1, keepdim=True), 0, 1) 19 | output_for_rknn.append(cls_sum) 20 | ``` 21 | 22 | (v6n, v6s) 可以在 yolov6/models/heads/effidehead_distill_ns.py 78~94行位置，注释掉这部分 23 | 24 | ``` 25 | cls_sum = torch.clamp(y[-1].sum(1, keepdim=True), 0, 1) 26 | output_for_rknn.append(cls_sum) 27 | ``` 28 | 29 | 30 | 31 | 32 | - 以上优化只影响了模型的导出，不影响训练过程，**训练步骤请参考 YOLOv6 官方文档**。 33 | 34 | 35 | 36 | 37 | ### 2.导出模型操作 38 | 39 | 在满足 ./requirements.txt 的环境要求后，执行以下语句导出模型 40 | 41 | ``` 42 | python deploy/RKNN/export_onnx_for_rknn.py --weight ./yolov6n.pt 43 | 44 | # 如果自己训练模型，则路径./yolov6n.pt 请改为自己模型的路径 45 | ``` 46 | 47 | 48 | 49 | ### 3.转RKNN模型、Python demo、C demo 50 | 51 | 请参考 https://github.com/airockchip/rknn_model_zoo/tree/main/models/CV/object_detection/yolo -------------------------------------------------------------------------------- /deploy/RKNN/export_onnx_for_rknn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import argparse 4 | import time 5 | import sys 6 | import os 7 | import torch 8 | import torch.nn as nn 9 | import onnx 10 | 11 | ROOT = os.getcwd() 12 | if str(ROOT) not in sys.path: 13 | sys.path.append(str(ROOT)) 14 | 15 | from yolov6.models.yolo import * 16 | from yolov6.models.effidehead import Detect 17 | from yolov6.layers.common import * 18 | from yolov6.utils.events import LOGGER 19 | from yolov6.utils.checkpoint import load_checkpoint 20 | 21 | 22 | if __name__ == '__main__': 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--weights', type=str, default='./yolov6s6.pt', help='weights path') 25 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width 26 | parser.add_argument('--batch-size', type=int, default=1, help='batch size') 27 | parser.add_argument('--half', action='store_true', help='FP16 half-precision export') 28 | parser.add_argument('--inplace', action='store_true', help='set Detect() inplace=True') 29 | parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 30 | args = parser.parse_args() 31 | args.img_size *= 2 if len(args.img_size) == 1 else 1 # expand 32 | print(args) 33 | t = time.time() 34 | 35 | # Check device 36 | cuda = args.device != 'cpu' and torch.cuda.is_available() 37 | device = torch.device('cuda:0' if cuda else 'cpu') 38 | assert not (device.type == 'cpu' and args.half), '--half only compatible with GPU export, i.e. use --device 0' 39 | # Load PyTorch model 40 | model = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True) # load FP32 model 41 | for layer in model.modules(): 42 | if isinstance(layer, RepVGGBlock): 43 | layer.switch_to_deploy() 44 | 45 | # Input 46 | img = torch.zeros(args.batch_size, 3, *args.img_size).to(device) # image size(1,3,320,192) iDetection 47 | 48 | # Update model 49 | if args.half: 50 | img, model = img.half(), model.half() # to FP16 51 | model.eval() 52 | for k, m in model.named_modules(): 53 | if isinstance(m, Conv): # assign export-friendly activations 54 | if isinstance(m.act, nn.SiLU): 55 | m.act = SiLU() 56 | elif isinstance(m, Detect): 57 | m.inplace = args.inplace 58 | 59 | model.detect.export_rknn = True 60 | y = model(img) # dry run 61 | 62 | # ONNX export 63 | try: 64 | LOGGER.info('\nStarting to export ONNX with rknn-optimized...') 65 | export_file = args.weights.replace('.pt', '.onnx') # filename 66 | torch.onnx.export(model, img, export_file, verbose=False, opset_version=12, 67 | training=torch.onnx.TrainingMode.EVAL, 68 | do_constant_folding=True, 69 | input_names=['image_arrays'], 70 | output_names=['outputs'], 71 | ) 72 | 73 | # Checks 74 | onnx_model = onnx.load(export_file) # load onnx model 75 | onnx.checker.check_model(onnx_model) # check onnx model 76 | LOGGER.info(f'ONNX with rknn-optimized export success, saved as {export_file}') 77 | except Exception as e: 78 | LOGGER.info(f'ONNX export failure: {e}') -------------------------------------------------------------------------------- /deploy/TensorRT/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | 3 | project(yolov6) 4 | 5 | add_definitions(-std=c++11) 6 | 7 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 8 | set(CMAKE_CXX_STANDARD 11) 9 | set(CMAKE_BUILD_TYPE Debug) 10 | # add_definitions("-Wall -g") 11 | find_package(CUDA REQUIRED) 12 | 13 | include_directories(${PROJECT_SOURCE_DIR}/include) 14 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 15 | # cuda 16 | include_directories(/usr/local/cuda/include) 17 | link_directories(/usr/local/cuda/lib64) 18 | # tensorrt 19 | include_directories(/usr/include/x86_64-linux-gnu/) 20 | link_directories(/usr/lib/x86_64-linux-gnu/) 21 | 22 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") 23 | 24 | find_package(OpenCV) 25 | include_directories(${OpenCV_INCLUDE_DIRS}) 26 | 27 | add_executable(yolov6 ${PROJECT_SOURCE_DIR}/yolov6.cpp) 28 | target_link_libraries(yolov6 nvinfer) 29 | target_link_libraries(yolov6 cudart) 30 | target_link_libraries(yolov6 ${OpenCV_LIBS}) 31 | 32 | add_definitions(-O2 -pthread) 33 | -------------------------------------------------------------------------------- /deploy/TensorRT/README.md: -------------------------------------------------------------------------------- 1 | # YOLOv6-TensorRT in C++ 2 | 3 | ## Dependencies 4 | - TensorRT-8.2.3.0 5 | - OpenCV-4.1.0 6 | 7 | 8 | 9 | ## Step 1: Get onnx model 10 | 11 | Follow the file [ONNX README](../../tools/quantization/tensorrt/post_training/README.md) to convert the pt model to onnx `yolov6n.onnx`. 12 | **Now don't support end2end onnx model which include the nms plugin** 13 | ```shell 14 | python ./deploy/ONNX/export_onnx.py \ 15 | --weights yolov6n.pt \ 16 | --img 640 \ 17 | --batch 1 18 | ``` 19 | 20 | ## Step 2: Prepare serialized engine file 21 | 22 | Follow the file [post training README](../../tools/quantization/tensorrt/post_training/README.md) to convert and save the serialized engine file `yolov6.engine`. 23 | 24 | ```shell 25 | python3 onnx_to_tensorrt.py --model ${ONNX_MODEL} \ 26 | --dtype int8 \ 27 | --max_calibration_size=${MAX_CALIBRATION_SIZE} \ 28 | --calibration-data=${CALIBRATION_DATA} \ 29 | --calibration-cache=${CACHE_FILENAME} \ 30 | --preprocess_func=${PREPROCESS_FUNC} \ 31 | --explicit-batch \ 32 | --verbose 33 | 34 | ``` 35 | 36 | ## Step 3: build the demo 37 | 38 | Please follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) to install TensorRT. 39 | 40 | And you should set the TensorRT path and CUDA path in CMakeLists.txt. 41 | 42 | If you train your custom dataset, you may need to modify the value of `num_class, image width height, and class name`. 43 | 44 | ```c++ 45 | const int num_class = 80; 46 | static const int INPUT_W = 640; 47 | static const int INPUT_H = 640; 48 | static const char* class_names[] = { 49 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 50 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 51 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 52 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 53 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 54 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 55 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 56 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 57 | "hair drier", "toothbrush" 58 | }; 59 | ``` 60 | 61 | build the demo: 62 | 63 | ```shell 64 | mkdir build 65 | cd build 66 | cmake .. 67 | make 68 | ``` 69 | 70 | Then run the demo: 71 | 72 | ```shell 73 | ./yolov6 ../you.engine -i image_path 74 | ``` 75 | 76 | # Evaluate the performance 77 | You can evaluate the performance of the TensorRT model. 78 | ``` 79 | python deploy/TensorRT/eval_yolo_trt.py \ 80 | --imgs_dir /path/to/images/val \ 81 | --labels_dir /path/to/labels/val\ 82 | --annotations /path/to/coco/format/annotation/file \ --batch 1 \ 83 | --img_size 640 \ 84 | --model /path/to/tensorrt/model \ 85 | --do_pr_metric --is_coco 86 | ``` 87 | Tips: 88 | `--is_coco`: if you are evaluating the COCO dataset, add this, if not, do not add this parameter. 89 | `--do_pr_metric`: If you want to get PR metric, add this. 90 | 91 | For example: 92 | ``` 93 | python deploy/TensorRT/eval_yolo_trt.py \ 94 | --imgs_dir /workdir/datasets/coco/images/val2017/ \ 95 | --labels_dir /workdir/datasets/coco/labels/val2017\ 96 | --annotations /workdir/datasets/coco/annotations/instances_val2017.json \ 97 | --batch 1 \ 98 | --img_size 640 \ 99 | --model weights/yolov6n.trt \ 100 | --do_pr_metric --is_coco 101 | 102 | ``` 103 | -------------------------------------------------------------------------------- /deploy/TensorRT/calibrator.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tensorrt as trt 3 | import pycuda.driver as cuda 4 | import pycuda.autoinit 5 | import numpy as np 6 | import cv2 7 | import glob 8 | from tensorrt_processor import letterbox 9 | 10 | import ctypes 11 | import logging 12 | logger = logging.getLogger(__name__) 13 | ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_char_p 14 | ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] 15 | 16 | 17 | """ 18 | There are 4 types calibrator in TensorRT. 19 | trt.IInt8LegacyCalibrator 20 | trt.IInt8EntropyCalibrator 21 | trt.IInt8EntropyCalibrator2 22 | trt.IInt8MinMaxCalibrator 23 | """ 24 | 25 | IMG_FORMATS = [".bmp", ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".dng", ".webp", ".mpo"] 26 | IMG_FORMATS.extend([f.upper() for f in IMG_FORMATS]) 27 | 28 | class Calibrator(trt.IInt8MinMaxCalibrator): 29 | def __init__(self, stream, cache_file=""): 30 | trt.IInt8MinMaxCalibrator.__init__(self) 31 | self.stream = stream 32 | self.d_input = cuda.mem_alloc(self.stream.calibration_data.nbytes) 33 | self.cache_file = cache_file 34 | stream.reset() 35 | 36 | def get_batch_size(self): 37 | return self.stream.batch_size 38 | 39 | def get_batch(self, names): 40 | print("######################") 41 | print(names) 42 | print("######################") 43 | batch = self.stream.next_batch() 44 | if not batch.size: 45 | return None 46 | 47 | cuda.memcpy_htod(self.d_input, batch) 48 | return [int(self.d_input)] 49 | 50 | def read_calibration_cache(self): 51 | # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None. 52 | if os.path.exists(self.cache_file): 53 | with open(self.cache_file, "rb") as f: 54 | logger.info("Using calibration cache to save time: {:}".format(self.cache_file)) 55 | return f.read() 56 | 57 | def write_calibration_cache(self, cache): 58 | with open(self.cache_file, "wb") as f: 59 | logger.info("Caching calibration data for future use: {:}".format(self.cache_file)) 60 | f.write(cache) 61 | 62 | 63 | def process_image(img_src, img_size, stride): 64 | '''Process image before image inference.''' 65 | image = letterbox(img_src, img_size, auto=False, return_int=True)[0] 66 | # Convert 67 | image = image.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB 68 | image = np.ascontiguousarray(image).astype(np.float32) 69 | image /= 255. # 0 - 255 to 0.0 - 1.0 70 | return image 71 | 72 | class DataLoader: 73 | def __init__(self, batch_size, batch_num, calib_img_dir, input_w, input_h): 74 | self.index = 0 75 | self.length = batch_num 76 | self.batch_size = batch_size 77 | self.input_h = input_h 78 | self.input_w = input_w 79 | # self.img_list = [i.strip() for i in open('calib.txt').readlines()] 80 | self.img_list = [os.path.join(calib_img_dir, x) for x in os.listdir(calib_img_dir) if os.path.splitext(x)[-1] in IMG_FORMATS] 81 | assert len(self.img_list) > self.batch_size * self.length, \ 82 | '{} must contains more than '.format(calib_img_dir) + str(self.batch_size * self.length) + ' images to calib' 83 | print('found all {} images to calib.'.format(len(self.img_list))) 84 | self.calibration_data = np.zeros((self.batch_size, 3, input_h, input_w), dtype=np.float32) 85 | 86 | def reset(self): 87 | self.index = 0 88 | 89 | def next_batch(self): 90 | if self.index < self.length: 91 | for i in range(self.batch_size): 92 | assert os.path.exists(self.img_list[i + self.index * self.batch_size]), f'{self.img_list[i + self.index * self.batch_size]} not found!!' 93 | img = cv2.imread(self.img_list[i + self.index * self.batch_size]) 94 | img = process_image(img, [self.input_h, self.input_w], 32) 95 | 96 | self.calibration_data[i] = img 97 | 98 | self.index += 1 99 | return np.ascontiguousarray(self.calibration_data, dtype=np.float32) 100 | else: 101 | return np.array([]) 102 | 103 | def __len__(self): 104 | return self.length 105 | -------------------------------------------------------------------------------- /deploy/TensorRT/visualize.py: -------------------------------------------------------------------------------- 1 | """visualize.py 2 | 3 | This script is for visualization of YOLO models. 4 | """ 5 | import os 6 | import sys 7 | import json 8 | import argparse 9 | import math 10 | 11 | import cv2 12 | import torch 13 | from tensorrt_processor import Processor 14 | from tqdm import tqdm 15 | 16 | 17 | def parse_args(): 18 | """Parse input arguments.""" 19 | desc = 'Visualization of YOLO TRT model' 20 | parser = argparse.ArgumentParser(description=desc) 21 | parser.add_argument( 22 | '--imgs-dir', type=str, default='./coco_images/', 23 | help='directory of to be visualized images ./coco_images/') 24 | parser.add_argument( 25 | '--visual-dir', type=str, default='./visual_out', 26 | help='directory of visualized images ./visual_out') 27 | parser.add_argument('--batch-size', type=int, 28 | default=1, help='batch size for training: default 64') 29 | parser.add_argument( 30 | '-c', '--category-num', type=int, default=80, 31 | help='number of object categories [80]') 32 | parser.add_argument( 33 | '--img-size', nargs='+', type=int, default=[640, 640], help='image size') 34 | parser.add_argument( 35 | '-m', '--model', type=str, default='./weights/yolov5s-simple.trt', 36 | help=('trt model path')) 37 | parser.add_argument( 38 | '--conf-thres', type=float, default=0.03, 39 | help='object confidence threshold') 40 | parser.add_argument( 41 | '--iou-thres', type=float, default=0.65, 42 | help='IOU threshold for NMS') 43 | parser.add_argument('--shrink_size', type=int, default=6, help='load img with size (img_size - shrink_size), for better performace.') 44 | args = parser.parse_args() 45 | return args 46 | 47 | 48 | def check_args(args): 49 | """Check and make sure command-line arguments are valid.""" 50 | if not os.path.isdir(args.imgs_dir): 51 | sys.exit('%s is not a valid directory' % args.imgs_dir) 52 | if not os.path.exists(args.visual_dir): 53 | print("Directory {} does not exist, create it".format(args.visual_dir)) 54 | os.makedirs(args.visual_dir) 55 | 56 | 57 | def generate_results(processor, imgs_dir, visual_dir, jpgs, conf_thres, iou_thres, 58 | batch_size=1, img_size=[640,640], shrink_size=0): 59 | """Run detection on each jpg and write results to file.""" 60 | results = [] 61 | # pbar = tqdm(jpgs, desc="TRT-Model test in val datasets.") 62 | pbar = tqdm(range(math.ceil(len(jpgs) / batch_size)), desc="TRT-Model test in val datasets.") 63 | idx = 0 64 | num_visualized = 0 65 | for _ in pbar: 66 | imgs = torch.randn((batch_size, 3, 640, 640), dtype=torch.float32, device=torch.device('cuda:0')) 67 | source_imgs = [] 68 | image_names = [] 69 | shapes = [] 70 | for i in range(batch_size): 71 | if (idx == len(jpgs)): break 72 | img = cv2.imread(os.path.join(imgs_dir, jpgs[idx])) 73 | img_src = img.copy() 74 | # shapes.append(img.shape) 75 | h0, w0 = img.shape[:2] 76 | r = (max(img_size) - shrink_size) / max(h0, w0) 77 | if r != 1: 78 | img = cv2.resize( 79 | img, 80 | (int(w0 * r), int(h0 * r)), 81 | interpolation=cv2.INTER_AREA 82 | if r < 1 else cv2.INTER_LINEAR, 83 | ) 84 | h, w = img.shape[:2] 85 | imgs[i], pad = processor.pre_process(img) 86 | source_imgs.append(img_src) 87 | shape = (h0, w0), ((h / h0, w / w0), pad) 88 | shapes.append(shape) 89 | image_names.append(jpgs[idx]) 90 | idx += 1 91 | output = processor.inference(imgs) 92 | 93 | for j in range(len(shapes)): 94 | pred = processor.post_process(output[j].unsqueeze(0), shapes[j], conf_thres=conf_thres, iou_thres=iou_thres) 95 | image = source_imgs[j] 96 | for p in pred: 97 | x = float(p[0]) 98 | y = float(p[1]) 99 | w = float(p[2] - p[0]) 100 | h = float(p[3] - p[1]) 101 | s = float(p[4]) 102 | 103 | cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (255, 0, 0), 1) 104 | 105 | # print("saving to {}".format(os.path.join(visual_dir, image_names[j]))) 106 | cv2.imwrite("{}".format(os.path.join(visual_dir, image_names[j])), image) 107 | 108 | def main(): 109 | args = parse_args() 110 | check_args(args) 111 | 112 | assert args.model.endswith('.trt'), "Only support trt engine test" 113 | 114 | # setup processor 115 | processor = Processor(model=args.model) 116 | jpgs = [j for j in os.listdir(args.imgs_dir) if j.endswith('.jpg')] 117 | generate_results(processor, args.imgs_dir, args.visual_dir, jpgs, args.conf_thres, args.iou_thres, 118 | batch_size=args.batch_size, img_size = args.img_size, shrink_size=args.shrink_size) 119 | 120 | 121 | if __name__ == '__main__': 122 | main() 123 | -------------------------------------------------------------------------------- /docs/Test_speed.md: -------------------------------------------------------------------------------- 1 | # Test speed 2 | 3 | This guidence explains how to reproduce speed results of YOLOv6. For fair comparison, the speed results do not contain the time cost of data pre-processing and NMS post-processing. 4 | 5 | ## 0. Prepare model 6 | 7 | Download the models you want to test from the latest release. 8 | 9 | ## 1. Prepare testing environment 10 | 11 | Refer to README, install packages corresponding to CUDA, CUDNN and TensorRT version. 12 | 13 | Here, we use Torch1.8.0 inference on V100 and TensorRT 7.2 Cuda 10.2 Cudnn 8.0.2 on T4. 14 | 15 | ## 2. Reproduce speed 16 | 17 | #### 2.1 Torch Inference on V100 18 | 19 | To get inference speed without TensorRT on V100, you can run the following command: 20 | 21 | ```shell 22 | python tools/eval.py --data data/coco.yaml --batch 32 --weights yolov6n.pt --task speed [--half] 23 | ``` 24 | 25 | - Speed results with batchsize = 1 are unstable in multiple runs, thus we do not provide the bs1 speed results. 26 | 27 | #### 2.2 TensorRT Inference on T4 28 | 29 | To get inference speed with TensorRT in FP16 mode on T4, you can follow the steps below: 30 | 31 | First, export pytorch model as onnx format using the following command: 32 | 33 | ```shell 34 | python deploy/ONNX/export_onnx.py --weights yolov6n.pt --device 0 --simplify --batch [1 or 32] 35 | ``` 36 | 37 | Second, generate an inference trt engine and test speed using `trtexec`: 38 | 39 | ``` 40 | trtexec --explicitBatch --fp16 --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --buildOnly --workspace=1024 --onnx=yolov6n.onnx --saveEngine=yolov6n.trt 41 | 42 | trtexec --fp16 --avgRuns=1000 --workspace=1024 --loadEngine=yolov6n.trt 43 | ``` 44 | -------------------------------------------------------------------------------- /docs/Train_coco_data.md: -------------------------------------------------------------------------------- 1 | # Train COCO Dataset 2 | 3 | This guidence shows the training commands for reproducing our results on COCO Dataset. 4 | 5 | ## For P5 models 6 | 7 | #### YOLOv6-N 8 | 9 | ```shell 10 | # Step 1: Training a base model 11 | # Be sure to open use_dfl mode in config file (use_dfl=True, reg_max=16) 12 | 13 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 14 | --batch 128 \ 15 | --conf configs/yolov6n.py \ 16 | --data data/coco.yaml \ 17 | --epoch 300 \ 18 | --fuse_ab \ 19 | --device 0,1,2,3,4,5,6,7 \ 20 | --name yolov6n_coco 21 | 22 | # Step 2: Self-distillation training 23 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 24 | --batch 128 \ 25 | --conf configs/yolov6n.py \ 26 | --data data/coco.yaml \ 27 | --epoch 300 \ 28 | --device 0,1,2,3,4,5,6,7 \ 29 | --distill \ 30 | --teacher_model_path runs/train/yolov6n_coco/weights/best_ckpt.pt \ 31 | --name yolov6n_coco 32 | ``` 33 | 34 | 35 | #### YOLOv6-S/M/L 36 | 37 | ```shell 38 | # Step 1: Training a base model 39 | # Be sure to open use_dfl mode in config file (use_dfl=True, reg_max=16) 40 | 41 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 42 | --batch 256 \ 43 | --conf configs/yolov6s.py \ # yolov6m/yolov6l 44 | --data data/coco.yaml \ 45 | --epoch 300 \ 46 | --fuse_ab \ 47 | --device 0,1,2,3,4,5,6,7 \ 48 | --name yolov6s_coco # yolov6m_coco/yolov6l_coco 49 | 50 | # Step 2: Self-distillation training 51 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 52 | --batch 256 \ # 128 for distillation of yolov6l 53 | --conf configs/yolov6s.py \ # yolov6m/yolov6l 54 | --data data/coco.yaml \ 55 | --epoch 300 \ 56 | --device 0,1,2,3,4,5,6,7 \ 57 | --distill \ 58 | --teacher_model_path runs/train/yolov6s_coco/weights/best_ckpt.pt \ 59 | --name yolov6s_coco # yolov6m_coco/yolov6l_coco 60 | 61 | ``` 62 | 63 | ## For P6 models 64 | 65 | #### YOLOv6-N6/S6 66 | 67 | ```shell 68 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 69 | --batch 128 \ 70 | --img 1280 \ 71 | --conf configs/yolov6s6.py \ # yolov6n6 72 | --data data/coco.yaml \ 73 | --epoch 300 \ 74 | --bs_per_gpu 16 \ 75 | --device 0,1,2,3,4,5,6,7 \ 76 | --name yolov6s6_coco # yolov6n6_coco 77 | 78 | ``` 79 | 80 | 81 | #### YOLOv6-M6/L6 82 | 83 | ```shell 84 | # Step 1: Training a base model 85 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 86 | --batch 128 \ 87 | --conf configs/yolov6l6.py \ # yolov6m6 88 | --data data/coco.yaml \ 89 | --epoch 300 \ 90 | --bs_per_gpu 16 \ 91 | --device 0,1,2,3,4,5,6,7 \ 92 | --name yolov6l6_coco # yolov6m6_coco 93 | 94 | # Step 2: Self-distillation training 95 | python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \ 96 | --batch 128 \ 97 | --conf configs/yolov6l6.py \ # yolov6m6 98 | --data data/coco.yaml \ 99 | --epoch 300 \ 100 | --bs_per_gpu 16 \ 101 | --device 0,1,2,3,4,5,6,7 \ 102 | --distill \ 103 | --teacher_model_path runs/train/yolov6l6_coco/weights/best_ckpt.pt \ 104 | --name yolov6l6_coco # yolov6m6_coco 105 | 106 | ``` 107 | -------------------------------------------------------------------------------- /docs/Tutorial of Quantization.md: -------------------------------------------------------------------------------- 1 | # Quantization Practice for YOLOv6 2 | For industrial deployment, it has been common practice to adopt quantization to further speed up runtime without much performance compromise. However, due to the heavy use of re-parameterization blocks in YOLOv6, previous PTQ techniques fail to produce high performance, while it is hard to incorporate QAT when it comes to matching fake quantizers during training and inference. 3 | 4 | In order to solve the quantization problem of YOLOv6, we firstly reconstruct the network with RepOptimizer, and then perform well-designed PTQ and QAT skills on this model. Finally we can obtain a SOTA quantized result(mAP 43.3 at 869 QPS) for YOLOv6s. 5 | 6 | Specific tutorials, please refer to the following links: 7 | * [Tutorial of RepOpt for YOLOv6](./tutorial_repopt.md) 8 | * [Tutorial of QAT for YOLOv6](../tools/qat/README.md) 9 | * [Partial Quantization](../tools/partial_quantization) 10 | * [PPQ Quantization](../tools/quantization/ppq) 11 | -------------------------------------------------------------------------------- /docs/tutorial_repopt.md: -------------------------------------------------------------------------------- 1 | # RepOpt version implementation of YOLOv6 2 | ## Introduction 3 | This is a RepOpt-version implementation of YOLOv6 according to RepOptimizer: https://arxiv.org/pdf/2205.15242.pdf @DingXiaoH \ 4 | It shows some advantages： 5 | 1. With only minor changes. it is compatible with the original repvgg version, and it is easy to reproduce the precision comparable with original version. 6 | 2. No more train/deploy transform. The target network is consistent when training and deploying. 7 | 3. A slight training acceleration of about 8%. 8 | 4. Last and the most important, It is quantization friendly. Compared to the original version, the mAP decrease of PTQ can be greatly improved. Furthermore, the architecture of RepOptimizer is friendly to wrap quant-models for QAT. 9 | 10 | ## Training 11 | The training of V6-RepOpt can be divided into two stages, hyperparameter search and target network training. 12 | 1. hyperparameter search. This stage is used to get a suitable 'scale' for RepOptimizer, and the result checkpoint can be passed to stage2. Remember to add `training_mode='hyper_search'` in your config. 13 | ``` 14 | python tools/train.py --batch 32 --conf configs/repopt/yolov6s_hs.py --data data/coco.yaml --device 0 15 | ``` 16 | Or you can directly use the [pretrained scale](https://github.com/xingyueye/YOLOv6/releases/download/0.1.0/yolov6s_scale.pt) we provided and omit this stage. 17 | 18 | 2. Training. Add the flag of `training_mode='repopt'` and pretraind model `scales='./assets/yolov6s_scale.pt',` in your config 19 | ``` 20 | python tools/train.py --batch 32 --conf configs/repopt/yolov6s_opt.py --data data/coco.yaml --device 0 21 | ``` 22 | ## Evaluation 23 | Reproduce mAP on COCO val2017 dataset, you can directly test our [pretrained model](https://github.com/xingyueye/YOLOv6/releases/download/0.1.0/yolov6s_opt.pt). 24 | ``` 25 | python tools/eval.py --data data/coco.yaml --batch 32 --weights yolov6s_opt.pt --task val 26 | ``` 27 | ## Benchmark 28 | We train a yolov6s-repopt with 300epochs, the fp32 mAP is 42.4, while the mAP of PTQ is 40.5. More results is coming soon... 29 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | # python3.8 environment 3 | 4 | torch>=1.8.0 5 | torchvision>=0.9.0 6 | numpy>=1.24.0 7 | opencv-python>=4.1.2 8 | PyYAML>=5.3.1 9 | scipy>=1.4.1 10 | tqdm>=4.41.0 11 | addict>=2.4.0 12 | tensorboard>=2.7.0 13 | pycocotools>=2.0 14 | onnx>=1.10.0 # ONNX export 15 | onnx-simplifier>=0.3.6 # ONNX simplifier 16 | thop # FLOPs computation 17 | # pytorch_quantization>=2.1.1 18 | -------------------------------------------------------------------------------- /tools/partial_quantization/README.md: -------------------------------------------------------------------------------- 1 | # Partial Quantization 2 | The performance of YOLOv6s heavily degrades from 42.4% to 35.6% after traditional PTQ, which is unacceptable. To resolve this issue, we propose **partial quantization**. First we analyze the quantization sensitivity of all layers, and then we let the most sensitive layers to have full precision as a compromise. 3 | 4 | With partial quantization, we finally reach 42.1%, only 0.3% loss in accuracy, while the throughput of the partially quantized model is about 1.56 times that of the FP16 model at a batch size of 32. This method achieves a nice tradeoff between accuracy and throughput. 5 | 6 | ## Prerequirements 7 | ```python 8 | pip install --extra-index-url=https://pypi.ngc.nvidia.com --trusted-host pypi.ngc.nvidia.com nvidia-pyindex 9 | pip install --extra-index-url=https://pypi.ngc.nvidia.com --trusted-host pypi.ngc.nvidia.com pytorch_quantization 10 | ``` 11 | ## Sensitivity analysis 12 | 13 | Please use the following command to perform sensitivity analysis. Since we randomly sample 128 images from train dataset each time, the sensitivity files will be slightly different. 14 | 15 | ```python 16 | python3 sensitivity_analyse.py --weights yolov6s_reopt.pt \ 17 | --batch-size 32 \ 18 | --batch-number 4 \ 19 | --data-root train_data_path 20 | ``` 21 | 22 | ## Partial quantization 23 | 24 | With the sensitivity file at hand, we then proceed with partial quantization as follows. 25 | 26 | ```python 27 | python3 partial_quant.py --weights yolov6s_reopt.pt \ 28 | --calib-weights yolov6s_repot_calib.pt \ 29 | --sensitivity-file yolov6s_reopt_sensivitiy_128_calib.txt \ 30 | --quant-boundary 55 \ 31 | --export-batch-size 1 32 | ``` 33 | 34 | ## Deployment 35 | 36 | Build a TRT engine 37 | 38 | ```python 39 | trtexec --workspace=1024 --percentile=99 --streams=1 --int8 --fp16 --avgRuns=10 --onnx=yolov6s_reopt_partial_bs1.sim.onnx --saveEngine=yolov6s_reopt_partial_bs1.sim.trt 40 | ``` 41 | 42 | ## Performance 43 | | Model | Size | Precision |mAPval
0.5:0.95 | SpeedT4
trt b1
(fps) | SpeedT4
trt b32
(fps) | 44 | | :-------------- | ----------- | ----------- |:----------------------- | ---------------------------------------- | -----------------------------------| 45 | | [**YOLOv6-s-partial**]
[bs1](https://github.com/lippman1125/YOLOv6/releases/download/0.1.0/yolov6s_reopt_partial_bs1.sim.onnx)
[bs32](https://github.com/lippman1125/YOLOv6/releases/download/0.1.0/yolov6s_reopt_partial_bs32.sim.onnx)
| 640 | INT8 |42.1 | 503 | 811 | 46 | | [**YOLOv6-s**] | 640 | FP16 |42.4 | 373 | 520 | 47 | -------------------------------------------------------------------------------- /tools/partial_quantization/eval.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from yolov6.core.evaler import Evaler 4 | 5 | class EvalerWrapper(object): 6 | def __init__(self, eval_cfg): 7 | task = eval_cfg['task'] 8 | save_dir = eval_cfg['save_dir'] 9 | half = eval_cfg['half'] 10 | data = eval_cfg['data'] 11 | batch_size = eval_cfg['batch_size'] 12 | img_size = eval_cfg['img_size'] 13 | device = eval_cfg['device'] 14 | dataloader = None 15 | 16 | Evaler.check_task(task) 17 | if not os.path.exists(save_dir): 18 | os.makedirs(save_dir) 19 | 20 | # reload thres/device/half/data according task 21 | conf_thres = 0.03 22 | iou_thres = 0.65 23 | device = Evaler.reload_device(device, None, task) 24 | data = Evaler.reload_dataset(data) if isinstance(data, str) else data 25 | 26 | # init 27 | val = Evaler(data, batch_size, img_size, conf_thres, \ 28 | iou_thres, device, half, save_dir) 29 | val.stride = eval_cfg['stride'] 30 | dataloader = val.init_data(dataloader, task) 31 | 32 | self.eval_cfg = eval_cfg 33 | self.half = half 34 | self.device = device 35 | self.task = task 36 | self.val = val 37 | self.val_loader = dataloader 38 | 39 | def eval(self, model): 40 | model.eval() 41 | model.to(self.device) 42 | if self.half is True: 43 | model.half() 44 | 45 | with torch.no_grad(): 46 | pred_result, vis_outputs, vis_paths = self.val.predict_model(model, self.val_loader, self.task) 47 | eval_result = self.val.eval_model(pred_result, model, self.val_loader, self.task) 48 | 49 | return eval_result 50 | -------------------------------------------------------------------------------- /tools/partial_quantization/eval.yaml: -------------------------------------------------------------------------------- 1 | task: 'val' 2 | save_dir: 'runs/val/exp' 3 | half: False 4 | data: '../../data/coco.yaml' 5 | batch_size: 32 6 | img_size: 640 7 | device: '0' 8 | stride: 32 9 | -------------------------------------------------------------------------------- /tools/partial_quantization/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pytorch_quantization import nn as quant_nn 3 | 4 | 5 | def set_module(model, submodule_key, module): 6 | tokens = submodule_key.split('.') 7 | sub_tokens = tokens[:-1] 8 | cur_mod = model 9 | for s in sub_tokens: 10 | cur_mod = getattr(cur_mod, s) 11 | setattr(cur_mod, tokens[-1], module) 12 | 13 | 14 | def get_module(model, submodule_key): 15 | sub_tokens = submodule_key.split('.') 16 | cur_mod = model 17 | for s in sub_tokens: 18 | cur_mod = getattr(cur_mod, s) 19 | return cur_mod 20 | 21 | 22 | def module_quant_disable(model, k): 23 | cur_module = get_module(model, k) 24 | if hasattr(cur_module, '_input_quantizer'): 25 | cur_module._input_quantizer.disable() 26 | if hasattr(cur_module, '_weight_quantizer'): 27 | cur_module._weight_quantizer.disable() 28 | 29 | 30 | def module_quant_enable(model, k): 31 | cur_module = get_module(model, k) 32 | if hasattr(cur_module, '_input_quantizer'): 33 | cur_module._input_quantizer.enable() 34 | if hasattr(cur_module, '_weight_quantizer'): 35 | cur_module._weight_quantizer.enable() 36 | 37 | 38 | def model_quant_disable(model): 39 | for name, module in model.named_modules(): 40 | if isinstance(module, quant_nn.TensorQuantizer): 41 | module.disable() 42 | 43 | 44 | def model_quant_enable(model): 45 | for name, module in model.named_modules(): 46 | if isinstance(module, quant_nn.TensorQuantizer): 47 | module.enable() 48 | 49 | 50 | def concat_quant_amax_fuse(ops_list): 51 | if len(ops_list) <= 1: 52 | return 53 | 54 | amax = -1 55 | for op in ops_list: 56 | if hasattr(op, '_amax'): 57 | op_amax = op._amax.detach().item() 58 | elif hasattr(op, '_input_quantizer'): 59 | op_amax = op._input_quantizer._amax.detach().item() 60 | else: 61 | print("Not quantable op, skip") 62 | return 63 | print("op amax = {:7.4f}, amax = {:7.4f}".format(op_amax, amax)) 64 | if amax < op_amax: 65 | amax = op_amax 66 | 67 | print("amax = {:7.4f}".format(amax)) 68 | for op in ops_list: 69 | if hasattr(op, '_amax'): 70 | op._amax.fill_(amax) 71 | elif hasattr(op, '_input_quantizer'): 72 | op._input_quantizer._amax.fill_(amax) 73 | 74 | 75 | def quant_sensitivity_load(file): 76 | assert os.path.exists(file), print("File {} does not exist".format(file)) 77 | quant_sensitivity = list() 78 | with open(file, 'r') as qfile: 79 | lines = qfile.readlines() 80 | for line in lines: 81 | layer, mAP1, mAP2 = line.strip('\n').split(' ') 82 | quant_sensitivity.append((layer, float(mAP1), float(mAP2))) 83 | 84 | return quant_sensitivity 85 | 86 | 87 | def quant_sensitivity_save(quant_sensitivity, file): 88 | with open(file, 'w') as qfile: 89 | for item in quant_sensitivity: 90 | name, mAP1, mAP2 = item 91 | line = name + " " + "{:0.4f}".format(mAP1) + " " + "{:0.4f}".format(mAP2) + "\n" 92 | qfile.write(line) 93 | -------------------------------------------------------------------------------- /tools/qat/README.md: -------------------------------------------------------------------------------- 1 | # Quantization-Aware Training 2 | 3 | As of v0.2.0 release, traditional post-training quantization (PTQ) produces a degraded performance of `YOLOv6-S` from 43.4% to 41.2%. This is however much improved compared with v0.1.0 since the most sensitve layers are removed. Yet it is not ready for deployment. Meanwhile, due to the inconsistency of reparameterization blocks during training and inference, quantization-aware training (QAT) cannot be directly integrated into YOLOv6. As a remedy, we first train a single-branch network called `YOLOv6-S-RepOpt` with [RepOptimizer](https://arxiv.org/pdf/2205.15242.pdf). It reaches 43.1% mAP and is very close to YOLOv6-S. We then apply our quantization strategy on `YOLOv6-S-RepOpt`. 4 | 5 | We apply post-training quantization to `YOLOv6-S-RepOpt`, and its mAP slightly drops by 0.5%. Hence it is necessary to use QAT to further improve the accuracy. Besides, we involve **channel-wise distillation** to accelerate the convergence. We finally reach a quantized model at 43.0% mAP. 6 | 7 | To deploy the quantized model on typical NVIDIA GPUs (e.g. T4), we export the model to the ONNX format, then we use TensorRT to build a serialized engine along with the computed scale cache. The performance arrives at **43.3% mAP**, only 0.1% left to match the fully float precision of `YOLOv6-S`. 8 | 9 | 10 | ## Pre-requirements 11 | 12 | It is required to install `pytorch_quantization`, on top of which we build our quantization strategy. 13 | 14 | ```python 15 | pip install --extra-index-url=https://pypi.ngc.nvidia.com --trusted-host pypi.ngc.nvidia.com nvidia-pyindex 16 | pip install --extra-index-url=https://pypi.ngc.nvidia.com --trusted-host pypi.ngc.nvidia.com pytorch_quantization 17 | ``` 18 | 19 | ## Training with RepOptimizer 20 | Firstly, train a `YOLOv6-S RepOpt` as follows, or download our realeased [checkpoint](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6s_v2_reopt.pt) and [scales](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6s_v2_scale.pt). 21 | * [Tutorial of RepOpt for YOLOv6](https://github.com/meituan/YOLOv6/blob/main/docs/tutorial_repopt.md) 22 | ## PTQ 23 | We perform PTQ to get the range of activations and weights. 24 | ```python 25 | CUDA_VISIBLE_DEVICES=0 python tools/train.py \ 26 | --data ./data/coco.yaml \ 27 | --output-dir ./runs/opt_train_v6s_ptq \ 28 | --conf configs/repopt/yolov6s_opt_qat.py \ 29 | --quant \ 30 | --calib \ 31 | --batch 32 \ 32 | --workers 0 33 | ``` 34 | 35 | ## QAT 36 | 37 | Our proposed QAT strategy comes with channel-wise distillation. It loades calibrated ReOptimizer-trained model and trains for 10 epochs. To reproduce the result, 38 | 39 | ```python 40 | CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 \ 41 | tools/train.py \ 42 | --data ./data/coco.yaml \ 43 | --output-dir ./runs/opt_train_v6s_qat \ 44 | --conf configs/repopt/yolov6s_opt_qat.py \ 45 | --quant \ 46 | --distill \ 47 | --distill_feat \ 48 | --batch 128 \ 49 | --epochs 10 \ 50 | --workers 32 \ 51 | --teacher_model_path ./assets/yolov6s_v2_reopt_43.1.pt \ 52 | --device 0,1,2,3,4,5,6,7 53 | ``` 54 | ## ONNX Export 55 | To export to ONNX, 56 | ```python 57 | python3 qat_export.py --weights yolov6s_v2_reopt_43.1.pt --quant-weights yolov6s_v2_reopt_qat_43.0.pt --graph-opt --export-batch-size 1 58 | ``` 59 | 60 | ## TensorRT Deployment 61 | 62 | To build a TRT engine, 63 | 64 | ```python 65 | trtexec --workspace=1024 --percentile=99 --streams=1 --int8 --fp16 --avgRuns=10 --onnx=yolov6s_v2_reopt_qat_43.0_bs1.sim.onnx --calib=yolov6s_v2_reopt_qat_43.0_remove_qdq_bs1_calibration_addscale.cache --saveEngine=yolov6s_v2_reopt_qat_43.0_bs1.sim.trt 66 | ``` 67 | You can directly build engine with [yolov6s_v2_quant.onnx](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6s_v2_reopt_qat_43.0_remove_qdq_bs1.sim.onnx) and [yolov6s_v2_calibration.cache](https://github.com/meituan/YOLOv6/releases/download/0.2.0/yolov6s_v2_reopt_qat_43.0_remove_qdq_bs1_calibration_addscale.cache) 68 | 69 | ## Performance Comparison 70 | 71 | We release our quantized and graph-optimized YOLOv6-S (v0.2.0) model. The following throughput is tested with TensorRT 8.4 on a NVIDIA Tesla T4 GPU. 72 | 73 | | Model | Size | Precision |mAPval
0.5:0.95 | SpeedT4
trt b1
(fps) | SpeedT4
trt b32
(fps) | 74 | | :-------------- | ----------- | ----------- |:----------------------- | ---------------------------------------- | -----------------------------------| 75 | | [**YOLOv6-S RepOpt**] | 640 | INT8 |43.3 | 619 | 924 | 76 | | [**YOLOv6-S**] | 640 | FP16 |43.4 | 377 | 541 | 77 | | [**YOLOv6-T RepOpt**] | 640 | INT8 |39.8 | 741 | 1167 | 78 | | [**YOLOv6-T**] | 640 | FP16 |40.3 | 449 | 659 | 79 | | [**YOLOv6-N RepOpt**] | 640 | INT8 |34.8 | 1114 | 1828 | 80 | | [**YOLOv6-N**] | 640 | FP16 |35.9 | 802 | 1234 | 81 | -------------------------------------------------------------------------------- /tools/quantization/mnn/README.md: -------------------------------------------------------------------------------- 1 | # Coming soon 2 | -------------------------------------------------------------------------------- /tools/quantization/ppq/write_qparams_onnx2trt.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import argparse 4 | import tensorrt as trt 5 | 6 | TRT_LOGGER = trt.Logger() 7 | 8 | EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) 9 | 10 | def GiB(val): 11 | return val * 1 << 30 12 | 13 | def json_load(filename): 14 | with open(filename) as json_file: 15 | data = json.load(json_file) 16 | return data 17 | 18 | def setDynamicRange(network, json_file): 19 | """Sets ranges for network layers.""" 20 | quant_param_json = json_load(json_file) 21 | act_quant = quant_param_json["act_quant_info"] 22 | 23 | for i in range(network.num_inputs): 24 | input_tensor = network.get_input(i) 25 | if act_quant.__contains__(input_tensor.name): 26 | print(input_tensor.name) 27 | value = act_quant[input_tensor.name] 28 | tensor_max = abs(value) 29 | tensor_min = -abs(value) 30 | input_tensor.dynamic_range = (tensor_min, tensor_max) 31 | 32 | for i in range(network.num_layers): 33 | layer = network.get_layer(i) 34 | 35 | for output_index in range(layer.num_outputs): 36 | tensor = layer.get_output(output_index) 37 | 38 | if act_quant.__contains__(tensor.name): 39 | print("\033[1;32mWrite quantization parameters:%s\033[0m" % tensor.name) 40 | value = act_quant[tensor.name] 41 | tensor_max = abs(value) 42 | tensor_min = -abs(value) 43 | tensor.dynamic_range = (tensor_min, tensor_max) 44 | else: 45 | print("\033[1;31mNo quantization parameters are written: %s\033[0m" % tensor.name) 46 | 47 | 48 | def build_engine(onnx_file, json_file, engine_file): 49 | builder = trt.Builder(TRT_LOGGER) 50 | network = builder.create_network(EXPLICIT_BATCH) 51 | 52 | config = builder.create_builder_config() 53 | 54 | # If it is a dynamic onnx model , you need to add the following. 55 | # profile = builder.create_optimization_profile() 56 | # profile.set_shape("input_name", (batch, channels, min_h, min_w), (batch, channels, opt_h, opt_w), (batch, channels, max_h, max_w)) 57 | # config.add_optimization_profile(profile) 58 | 59 | 60 | parser = trt.OnnxParser(network, TRT_LOGGER) 61 | config.max_workspace_size = GiB(1) 62 | 63 | if not os.path.exists(onnx_file): 64 | quit('ONNX file {} not found'.format(onnx_file)) 65 | 66 | with open(onnx_file, 'rb') as model: 67 | if not parser.parse(model.read()): 68 | print('ERROR: Failed to parse the ONNX file.') 69 | for error in range(parser.num_errors): 70 | print(parser.get_error(error)) 71 | return None 72 | 73 | config.set_flag(trt.BuilderFlag.INT8) 74 | 75 | setDynamicRange(network, json_file) 76 | 77 | engine = builder.build_engine(network, config) 78 | 79 | with open(engine_file, "wb") as f: 80 | f.write(engine.serialize()) 81 | 82 | 83 | if __name__ == '__main__': 84 | # Add plugins if needed 85 | # import ctypes 86 | # ctypes.CDLL("libmmdeploy_tensorrt_ops.so") 87 | parser = argparse.ArgumentParser(description='Writing qparams to onnx to convert tensorrt engine.') 88 | parser.add_argument('--onnx', type=str, default=None) 89 | parser.add_argument('--qparam_json', type=str, default=None) 90 | parser.add_argument('--engine', type=str, default=None) 91 | arg = parser.parse_args() 92 | 93 | build_engine(arg.onnx, arg.qparam_json, arg.engine) 94 | print("\033[1;32mgenerate %s\033[0m" % arg.engine) 95 | 96 | 97 | -------------------------------------------------------------------------------- /tools/quantization/tensorrt/post_training/README.md: -------------------------------------------------------------------------------- 1 | # ONNX -> TensorRT INT8 2 | These scripts were last tested using the 3 | [NGC TensorRT Container Version 20.06-py3](https://ngc.nvidia.com/catalog/containers/nvidia:tensorrt). 4 | You can see the corresponding framework versions for this container [here](https://docs.nvidia.com/deeplearning/sdk/tensorrt-container-release-notes/rel_20.06.html#rel_20.06). 5 | 6 | ## Quickstart 7 | 8 | > **NOTE**: This INT8 example is only valid for **fixed-shape** ONNX models at the moment. 9 | > 10 | INT8 Calibration on **dynamic-shape** models is now supported, however this example has not been updated 11 | to reflect that yet. For more details on INT8 Calibration for **dynamic-shape** models, please 12 | see the [documentation](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#int8-calib-dynamic-shapes). 13 | 14 | ### 1. Convert ONNX model to TensorRT INT8 15 | 16 | See `./onnx_to_tensorrt.py -h` for full list of command line arguments. 17 | 18 | ```bash 19 | ./onnx_to_tensorrt.py --explicit-batch \ 20 | --onnx resnet50/model.onnx \ 21 | --fp16 \ 22 | --int8 \ 23 | --calibration-cache="caches/yolov6.cache" \ 24 | -o resnet50.int8.engine 25 | ``` 26 | 27 | See the [INT8 Calibration](#int8-calibration) section below for details on calibration 28 | using your own model or different data, where you don't have an existing calibration cache 29 | or want to create a new one. 30 | 31 | ## INT8 Calibration 32 | 33 | See [Calibrator.py](Calibrator.py) for a reference implementation 34 | of TensorRT's [IInt8EntropyCalibrator2](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/infer/Int8/EntropyCalibrator2.html). 35 | 36 | This class can be tweaked to work for other kinds of models, inputs, etc. 37 | 38 | In the [Quickstart](#quickstart) section above, we made use of a pre-existing cache, 39 | [caches/yolov6.cache](caches/yolov6.cache), to save time for the sake of an example. 40 | 41 | However, to calibrate using different data or a different model, you can do so with the `--calibration-data` argument. 42 | 43 | * This requires that you've mounted a dataset, such as Imagenet, to use for calibration. 44 | * Add something like `-v /imagenet:/imagenet` to your Docker command in Step (1) 45 | to mount a dataset found locally at `/imagenet`. 46 | * You can specify your own `preprocess_func` by defining it inside of `Calibrator.py` 47 | 48 | ```bash 49 | # Path to dataset to use for calibration. 50 | # **Not necessary if you already have a calibration cache from a previous run. 51 | CALIBRATION_DATA="/imagenet" 52 | 53 | # Truncate calibration images to a random sample of this amount if more are found. 54 | # **Not necessary if you already have a calibration cache from a previous run. 55 | MAX_CALIBRATION_SIZE=512 56 | 57 | # Calibration cache to be used instead of calibration data if it already exists, 58 | # or the cache will be created from the calibration data if it doesn't exist. 59 | CACHE_FILENAME="caches/yolov6.cache" 60 | 61 | # Path to ONNX model 62 | ONNX_MODEL="model/yolov6.onnx" 63 | 64 | # Path to write TensorRT engine to 65 | OUTPUT="yolov6.int8.engine" 66 | 67 | # Creates an int8 engine from your ONNX model, creating ${CACHE_FILENAME} based 68 | # on your ${CALIBRATION_DATA}, unless ${CACHE_FILENAME} already exists, then 69 | # it will use simply use that instead. 70 | python3 onnx_to_tensorrt.py --fp16 --int8 -v \ 71 | --max_calibration_size=${MAX_CALIBRATION_SIZE} \ 72 | --calibration-data=${CALIBRATION_DATA} \ 73 | --calibration-cache=${CACHE_FILENAME} \ 74 | --preprocess_func=${PREPROCESS_FUNC} \ 75 | --explicit-batch \ 76 | --onnx ${ONNX_MODEL} -o ${OUTPUT} 77 | 78 | ``` 79 | 80 | ### Pre-processing 81 | 82 | In order to calibrate your model correctly, you should `pre-process` your data the same way 83 | that you would during inference. 84 | -------------------------------------------------------------------------------- /tools/quantization/tensorrt/post_training/quant.sh: -------------------------------------------------------------------------------- 1 | # Path to ONNX model 2 | # ex: ../yolov6.onnx 3 | ONNX_MODEL=$1 4 | 5 | # Path to dataset to use for calibration. 6 | # **Not necessary if you already have a calibration cache from a previous run. 7 | CALIBRATION_DATA=$2 8 | 9 | # Path to Cache file to Serving 10 | # ex: ./caches/demo.cache 11 | CACHE_FILENAME=$3 12 | 13 | # Path to write TensorRT engine to 14 | OUTPUT=$4 15 | 16 | # Creates an int8 engine from your ONNX model, creating ${CACHE_FILENAME} based 17 | # on your ${CALIBRATION_DATA}, unless ${CACHE_FILENAME} already exists, then 18 | # it will use simply use that instead. 19 | python3 onnx_to_tensorrt.py --fp16 --int8 -v \ 20 | --calibration-data=${CALIBRATION_DATA} \ 21 | --calibration-cache=${CACHE_FILENAME} \ 22 | --explicit-batch \ 23 | --onnx ${ONNX_MODEL} -o ${OUTPUT} 24 | -------------------------------------------------------------------------------- /tools/quantization/tensorrt/requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | # python3.8 environment 3 | 4 | tensorrt # TensorRT 8.0+ 5 | pycuda==2020.1 # CUDA 11.0 6 | nvidia-pyindex 7 | pytorch-quantization 8 | -------------------------------------------------------------------------------- /tools/quantization/tensorrt/training_aware/QAT_quantizer.py: -------------------------------------------------------------------------------- 1 | # 2 | # QAT_quantizer.py 3 | # YOLOv6 4 | # 5 | # Created by Meituan on 2022/06/24. 6 | # Copyright © 2022 7 | # 8 | 9 | from absl import logging 10 | from pytorch_quantization import nn as quant_nn 11 | from pytorch_quantization import quant_modules 12 | 13 | # Call this function before defining the model 14 | def tensorrt_official_qat(): 15 | # Quantization Aware Training is based on Straight Through Estimator (STE) derivative approximation. 16 | # It is some time known as “quantization aware training”. 17 | 18 | # PyTorch-Quantization is a toolkit for training and evaluating PyTorch models with simulated quantization. 19 | # Quantization can be added to the model automatically, or manually, allowing the model to be tuned for accuracy and performance. 20 | # Quantization is compatible with NVIDIAs high performance integer kernels which leverage integer Tensor Cores. 21 | # The quantized model can be exported to ONNX and imported by TensorRT 8.0 and later. 22 | # https://github.com/NVIDIA/TensorRT/blob/main/tools/pytorch-quantization/examples/finetune_quant_resnet50.ipynb 23 | 24 | # The example to export the 25 | # model.eval() 26 | # quant_nn.TensorQuantizer.use_fb_fake_quant = True # We have to shift to pytorch's fake quant ops before exporting the model to ONNX 27 | # opset_version = 13 28 | 29 | # Export ONNX for multiple batch sizes 30 | # print("Creating ONNX file: " + onnx_filename) 31 | # dummy_input = torch.randn(batch_onnx, 3, 224, 224, device='cuda') #TODO: switch input dims by model 32 | # torch.onnx.export(model, dummy_input, onnx_filename, verbose=False, opset_version=opset_version, enable_onnx_checker=False, do_constant_folding=True) 33 | try: 34 | quant_modules.initialize() 35 | except NameError: 36 | logging.info("initialzation error for quant_modules") 37 | 38 | # def QAT_quantizer(): 39 | # coming soon 40 | -------------------------------------------------------------------------------- /yolov6/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/yolov6/__init__.py -------------------------------------------------------------------------------- /yolov6/assigners/__init__.py: -------------------------------------------------------------------------------- 1 | from .atss_assigner import ATSSAssigner 2 | from .tal_assigner import TaskAlignedAssigner 3 | -------------------------------------------------------------------------------- /yolov6/assigners/anchor_generator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def generate_anchors(feats, fpn_strides, grid_cell_size=5.0, grid_cell_offset=0.5, device='cpu', is_eval=False, mode='af'): 5 | '''Generate anchors from features.''' 6 | anchors = [] 7 | anchor_points = [] 8 | stride_tensor = [] 9 | num_anchors_list = [] 10 | assert feats is not None 11 | if is_eval: 12 | for i, stride in enumerate(fpn_strides): 13 | _, _, h, w = feats[i].shape 14 | shift_x = torch.arange(end=w, device=device) + grid_cell_offset 15 | shift_y = torch.arange(end=h, device=device) + grid_cell_offset 16 | shift_y, shift_x = torch.meshgrid(shift_y, shift_x) 17 | anchor_point = torch.stack( 18 | [shift_x, shift_y], axis=-1).to(torch.float) 19 | if mode == 'af': # anchor-free 20 | anchor_points.append(anchor_point.reshape([-1, 2])) 21 | stride_tensor.append( 22 | torch.full( 23 | (h * w, 1), stride, dtype=torch.float, device=device)) 24 | elif mode == 'ab': # anchor-based 25 | anchor_points.append(anchor_point.reshape([-1, 2]).repeat(3,1)) 26 | stride_tensor.append( 27 | torch.full( 28 | (h * w, 1), stride, dtype=torch.float, device=device).repeat(3,1)) 29 | anchor_points = torch.cat(anchor_points) 30 | stride_tensor = torch.cat(stride_tensor) 31 | return anchor_points, stride_tensor 32 | else: 33 | for i, stride in enumerate(fpn_strides): 34 | _, _, h, w = feats[i].shape 35 | cell_half_size = grid_cell_size * stride * 0.5 36 | shift_x = (torch.arange(end=w, device=device) + grid_cell_offset) * stride 37 | shift_y = (torch.arange(end=h, device=device) + grid_cell_offset) * stride 38 | shift_y, shift_x = torch.meshgrid(shift_y, shift_x) 39 | anchor = torch.stack( 40 | [ 41 | shift_x - cell_half_size, shift_y - cell_half_size, 42 | shift_x + cell_half_size, shift_y + cell_half_size 43 | ], 44 | axis=-1).clone().to(feats[0].dtype) 45 | anchor_point = torch.stack( 46 | [shift_x, shift_y], axis=-1).clone().to(feats[0].dtype) 47 | 48 | if mode == 'af': # anchor-free 49 | anchors.append(anchor.reshape([-1, 4])) 50 | anchor_points.append(anchor_point.reshape([-1, 2])) 51 | elif mode == 'ab': # anchor-based 52 | anchors.append(anchor.reshape([-1, 4]).repeat(3,1)) 53 | anchor_points.append(anchor_point.reshape([-1, 2]).repeat(3,1)) 54 | num_anchors_list.append(len(anchors[-1])) 55 | stride_tensor.append( 56 | torch.full( 57 | [num_anchors_list[-1], 1], stride, dtype=feats[0].dtype)) 58 | anchors = torch.cat(anchors) 59 | anchor_points = torch.cat(anchor_points).to(device) 60 | stride_tensor = torch.cat(stride_tensor).to(device) 61 | return anchors, anchor_points, num_anchors_list, stride_tensor 62 | 63 | -------------------------------------------------------------------------------- /yolov6/assigners/assigner_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | def dist_calculator(gt_bboxes, anchor_bboxes): 5 | """compute center distance between all bbox and gt 6 | 7 | Args: 8 | gt_bboxes (Tensor): shape(bs*n_max_boxes, 4) 9 | anchor_bboxes (Tensor): shape(num_total_anchors, 4) 10 | Return: 11 | distances (Tensor): shape(bs*n_max_boxes, num_total_anchors) 12 | ac_points (Tensor): shape(num_total_anchors, 2) 13 | """ 14 | gt_cx = (gt_bboxes[:, 0] + gt_bboxes[:, 2]) / 2.0 15 | gt_cy = (gt_bboxes[:, 1] + gt_bboxes[:, 3]) / 2.0 16 | gt_points = torch.stack([gt_cx, gt_cy], dim=1) 17 | ac_cx = (anchor_bboxes[:, 0] + anchor_bboxes[:, 2]) / 2.0 18 | ac_cy = (anchor_bboxes[:, 1] + anchor_bboxes[:, 3]) / 2.0 19 | ac_points = torch.stack([ac_cx, ac_cy], dim=1) 20 | 21 | distances = (gt_points[:, None, :] - ac_points[None, :, :]).pow(2).sum(-1).sqrt() 22 | 23 | return distances, ac_points 24 | 25 | def select_candidates_in_gts(xy_centers, gt_bboxes, eps=1e-9): 26 | """select the positive anchors's center in gt 27 | 28 | Args: 29 | xy_centers (Tensor): shape(bs*n_max_boxes, num_total_anchors, 4) 30 | gt_bboxes (Tensor): shape(bs, n_max_boxes, 4) 31 | Return: 32 | (Tensor): shape(bs, n_max_boxes, num_total_anchors) 33 | """ 34 | n_anchors = xy_centers.size(0) 35 | bs, n_max_boxes, _ = gt_bboxes.size() 36 | _gt_bboxes = gt_bboxes.reshape([-1, 4]) 37 | xy_centers = xy_centers.unsqueeze(0).repeat(bs * n_max_boxes, 1, 1) 38 | gt_bboxes_lt = _gt_bboxes[:, 0:2].unsqueeze(1).repeat(1, n_anchors, 1) 39 | gt_bboxes_rb = _gt_bboxes[:, 2:4].unsqueeze(1).repeat(1, n_anchors, 1) 40 | b_lt = xy_centers - gt_bboxes_lt 41 | b_rb = gt_bboxes_rb - xy_centers 42 | bbox_deltas = torch.cat([b_lt, b_rb], dim=-1) 43 | bbox_deltas = bbox_deltas.reshape([bs, n_max_boxes, n_anchors, -1]) 44 | return (bbox_deltas.min(axis=-1)[0] > eps).to(gt_bboxes.dtype) 45 | 46 | def select_highest_overlaps(mask_pos, overlaps, n_max_boxes): 47 | """if an anchor box is assigned to multiple gts, 48 | the one with the highest iou will be selected. 49 | 50 | Args: 51 | mask_pos (Tensor): shape(bs, n_max_boxes, num_total_anchors) 52 | overlaps (Tensor): shape(bs, n_max_boxes, num_total_anchors) 53 | Return: 54 | target_gt_idx (Tensor): shape(bs, num_total_anchors) 55 | fg_mask (Tensor): shape(bs, num_total_anchors) 56 | mask_pos (Tensor): shape(bs, n_max_boxes, num_total_anchors) 57 | """ 58 | fg_mask = mask_pos.sum(axis=-2) 59 | if fg_mask.max() > 1: 60 | mask_multi_gts = (fg_mask.unsqueeze(1) > 1).repeat([1, n_max_boxes, 1]) 61 | max_overlaps_idx = overlaps.argmax(axis=1) 62 | is_max_overlaps = F.one_hot(max_overlaps_idx, n_max_boxes) 63 | is_max_overlaps = is_max_overlaps.permute(0, 2, 1).to(overlaps.dtype) 64 | mask_pos = torch.where(mask_multi_gts, is_max_overlaps, mask_pos) 65 | fg_mask = mask_pos.sum(axis=-2) 66 | target_gt_idx = mask_pos.argmax(axis=-2) 67 | return target_gt_idx, fg_mask , mask_pos 68 | 69 | def iou_calculator(box1, box2, eps=1e-9): 70 | """Calculate iou for batch 71 | 72 | Args: 73 | box1 (Tensor): shape(bs, n_max_boxes, 1, 4) 74 | box2 (Tensor): shape(bs, 1, num_total_anchors, 4) 75 | Return: 76 | (Tensor): shape(bs, n_max_boxes, num_total_anchors) 77 | """ 78 | box1 = box1.unsqueeze(2) # [N, M1, 4] -> [N, M1, 1, 4] 79 | box2 = box2.unsqueeze(1) # [N, M2, 4] -> [N, 1, M2, 4] 80 | px1y1, px2y2 = box1[:, :, :, 0:2], box1[:, :, :, 2:4] 81 | gx1y1, gx2y2 = box2[:, :, :, 0:2], box2[:, :, :, 2:4] 82 | x1y1 = torch.maximum(px1y1, gx1y1) 83 | x2y2 = torch.minimum(px2y2, gx2y2) 84 | overlap = (x2y2 - x1y1).clip(0).prod(-1) 85 | area1 = (px2y2 - px1y1).clip(0).prod(-1) 86 | area2 = (gx2y2 - gx1y1).clip(0).prod(-1) 87 | union = area1 + area2 - overlap + eps 88 | 89 | return overlap / union 90 | -------------------------------------------------------------------------------- /yolov6/data/data_load.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # This code is based on 4 | # https://github.com/ultralytics/yolov5/blob/master/utils/dataloaders.py 5 | 6 | import os 7 | from torch.utils.data import dataloader, distributed 8 | 9 | from .datasets import TrainValDataset 10 | from yolov6.utils.events import LOGGER 11 | from yolov6.utils.torch_utils import torch_distributed_zero_first 12 | 13 | 14 | def create_dataloader( 15 | path, 16 | img_size, 17 | batch_size, 18 | stride, 19 | hyp=None, 20 | augment=False, 21 | check_images=False, 22 | check_labels=False, 23 | pad=0.0, 24 | rect=False, 25 | rank=-1, 26 | workers=8, 27 | shuffle=False, 28 | data_dict=None, 29 | task="Train", 30 | ): 31 | """Create general dataloader. 32 | 33 | Returns dataloader and dataset 34 | """ 35 | if rect and shuffle: 36 | LOGGER.warning( 37 | "WARNING: --rect is incompatible with DataLoader shuffle, setting shuffle=False" 38 | ) 39 | shuffle = False 40 | with torch_distributed_zero_first(rank): 41 | dataset = TrainValDataset( 42 | path, 43 | img_size, 44 | batch_size, 45 | augment=augment, 46 | hyp=hyp, 47 | rect=rect, 48 | check_images=check_images, 49 | check_labels=check_labels, 50 | stride=int(stride), 51 | pad=pad, 52 | rank=rank, 53 | data_dict=data_dict, 54 | task=task, 55 | ) 56 | 57 | batch_size = min(batch_size, len(dataset)) 58 | workers = min( 59 | [ 60 | os.cpu_count() // int(os.getenv("WORLD_SIZE", 1)), 61 | batch_size if batch_size > 1 else 0, 62 | workers, 63 | ] 64 | ) # number of workers 65 | sampler = ( 66 | None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) 67 | ) 68 | return ( 69 | TrainValDataLoader( 70 | dataset, 71 | batch_size=batch_size, 72 | shuffle=shuffle and sampler is None, 73 | num_workers=workers, 74 | sampler=sampler, 75 | pin_memory=True, 76 | collate_fn=TrainValDataset.collate_fn, 77 | ), 78 | dataset, 79 | ) 80 | 81 | 82 | class TrainValDataLoader(dataloader.DataLoader): 83 | """Dataloader that reuses workers 84 | 85 | Uses same syntax as vanilla DataLoader 86 | """ 87 | 88 | def __init__(self, *args, **kwargs): 89 | super().__init__(*args, **kwargs) 90 | object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler)) 91 | self.iterator = super().__iter__() 92 | 93 | def __len__(self): 94 | return len(self.batch_sampler.sampler) 95 | 96 | def __iter__(self): 97 | for i in range(len(self)): 98 | yield next(self.iterator) 99 | 100 | 101 | class _RepeatSampler: 102 | """Sampler that repeats forever 103 | 104 | Args: 105 | sampler (Sampler) 106 | """ 107 | 108 | def __init__(self, sampler): 109 | self.sampler = sampler 110 | 111 | def __iter__(self): 112 | while True: 113 | yield from iter(self.sampler) 114 | -------------------------------------------------------------------------------- /yolov6/data/vis_dataset.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Description: visualize yolo label image. 3 | 4 | import argparse 5 | import os 6 | import cv2 7 | import numpy as np 8 | 9 | IMG_FORMATS = ["bmp", "jpg", "jpeg", "png", "tif", "tiff", "dng", "webp", "mpo"] 10 | IMG_FORMATS.extend([f.upper() for f in IMG_FORMATS]) 11 | 12 | 13 | def main(args): 14 | img_dir, label_dir, class_names = args.img_dir, args.label_dir, args.class_names 15 | 16 | label_map = dict() 17 | for class_id, classname in enumerate(class_names): 18 | label_map[class_id] = classname 19 | 20 | for file in os.listdir(img_dir): 21 | if file.split('.')[-1] not in IMG_FORMATS: 22 | print(f'[Warning]: Non-image file {file}') 23 | continue 24 | img_path = os.path.join(img_dir, file) 25 | label_path = os.path.join(label_dir, file[: file.rindex('.')] + '.txt') 26 | 27 | try: 28 | img_data = cv2.imread(img_path) 29 | height, width, _ = img_data.shape 30 | color = [tuple(np.random.choice(range(256), size=3)) for i in class_names] 31 | thickness = 2 32 | 33 | with open(label_path, 'r') as f: 34 | for bbox in f: 35 | cls, x_c, y_c, w, h = [float(v) if i > 0 else int(v) for i, v in enumerate(bbox.split('\n')[0].split(' '))] 36 | 37 | x_tl = int((x_c - w / 2) * width) 38 | y_tl = int((y_c - h / 2) * height) 39 | cv2.rectangle(img_data, (x_tl, y_tl), (x_tl + int(w * width), y_tl + int(h * height)), tuple([int(x) for x in color[cls]]), thickness) 40 | cv2.putText(img_data, label_map[cls], (x_tl, y_tl - 10), cv2.FONT_HERSHEY_COMPLEX, 1, tuple([int(x) for x in color[cls]]), thickness) 41 | 42 | cv2.imshow('image', img_data) 43 | cv2.waitKey(0) 44 | except Exception as e: 45 | print(f'[Error]: {e} {img_path}') 46 | print('======All Done!======') 47 | 48 | 49 | if __name__ == '__main__': 50 | parser = argparse.ArgumentParser() 51 | parser.add_argument('--img_dir', default='VOCdevkit/voc_07_12/images') 52 | parser.add_argument('--label_dir', default='VOCdevkit/voc_07_12/labels') 53 | parser.add_argument('--class_names', default=['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 54 | 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']) 55 | 56 | args = parser.parse_args() 57 | print(args) 58 | 59 | main(args) 60 | -------------------------------------------------------------------------------- /yolov6/data/voc2yolo.py: -------------------------------------------------------------------------------- 1 | import xml.etree.ElementTree as ET 2 | from tqdm import tqdm 3 | import os 4 | import shutil 5 | import argparse 6 | 7 | # VOC dataset (refer https://github.com/ultralytics/yolov5/blob/master/data/VOC.yaml) 8 | # VOC2007 trainval: 446MB, 5012 images 9 | # VOC2007 test: 438MB, 4953 images 10 | # VOC2012 trainval: 1.95GB, 17126 images 11 | 12 | VOC_NAMES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 13 | 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'] 14 | 15 | 16 | def convert_label(path, lb_path, year, image_id): 17 | def convert_box(size, box): 18 | dw, dh = 1. / size[0], 1. / size[1] 19 | x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2] 20 | return x * dw, y * dh, w * dw, h * dh 21 | in_file = open(os.path.join(path, f'VOC{year}/Annotations/{image_id}.xml')) 22 | out_file = open(lb_path, 'w') 23 | tree = ET.parse(in_file) 24 | root = tree.getroot() 25 | size = root.find('size') 26 | w = int(size.find('width').text) 27 | h = int(size.find('height').text) 28 | for obj in root.iter('object'): 29 | cls = obj.find('name').text 30 | if cls in VOC_NAMES and not int(obj.find('difficult').text) == 1: 31 | xmlbox = obj.find('bndbox') 32 | bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')]) 33 | cls_id = VOC_NAMES.index(cls) # class id 34 | out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n') 35 | 36 | 37 | def gen_voc07_12(voc_path): 38 | ''' 39 | Generate voc07+12 setting dataset: 40 | train: # train images 16551 images 41 | - images/train2012 42 | - images/train2007 43 | - images/val2012 44 | - images/val2007 45 | val: # val images (relative to 'path') 4952 images 46 | - images/test2007 47 | ''' 48 | dataset_root = os.path.join(voc_path, 'voc_07_12') 49 | if not os.path.exists(dataset_root): 50 | os.makedirs(dataset_root) 51 | 52 | dataset_settings = {'train': ['train2007', 'val2007', 'train2012', 'val2012'], 'val':['test2007']} 53 | for item in ['images', 'labels']: 54 | for data_type, data_list in dataset_settings.items(): 55 | for data_name in data_list: 56 | ori_path = os.path.join(voc_path, item, data_name) 57 | new_path = os.path.join(dataset_root, item, data_type) 58 | if not os.path.exists(new_path): 59 | os.makedirs(new_path) 60 | 61 | print(f'[INFO]: Copying {ori_path} to {new_path}') 62 | for file in os.listdir(ori_path): 63 | shutil.copy(os.path.join(ori_path, file), new_path) 64 | 65 | 66 | def main(args): 67 | voc_path = args.voc_path 68 | for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'): 69 | imgs_path = os.path.join(voc_path, 'images', f'{image_set}') 70 | lbs_path = os.path.join(voc_path, 'labels', f'{image_set}') 71 | 72 | try: 73 | with open(os.path.join(voc_path, f'VOC{year}/ImageSets/Main/{image_set}.txt'), 'r') as f: 74 | image_ids = f.read().strip().split() 75 | if not os.path.exists(imgs_path): 76 | os.makedirs(imgs_path) 77 | if not os.path.exists(lbs_path): 78 | os.makedirs(lbs_path) 79 | 80 | for id in tqdm(image_ids, desc=f'{image_set}{year}'): 81 | f = os.path.join(voc_path, f'VOC{year}/JPEGImages/{id}.jpg') # old img path 82 | lb_path = os.path.join(lbs_path, f'{id}.txt') # new label path 83 | convert_label(voc_path, lb_path, year, id) # convert labels to YOLO format 84 | if os.path.exists(f): 85 | shutil.move(f, imgs_path) # move image 86 | except Exception as e: 87 | print(f'[Warning]: {e} {year}{image_set} convert fail!') 88 | 89 | gen_voc07_12(voc_path) 90 | 91 | 92 | 93 | if __name__ == '__main__': 94 | parser = argparse.ArgumentParser() 95 | parser.add_argument('--voc_path', default='VOCdevkit') 96 | 97 | args = parser.parse_args() 98 | print(args) 99 | 100 | main(args) 101 | -------------------------------------------------------------------------------- /yolov6/layers/dbb_transforms.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn.functional as F 4 | 5 | 6 | def transI_fusebn(kernel, bn): 7 | gamma = bn.weight 8 | std = (bn.running_var + bn.eps).sqrt() 9 | return kernel * ((gamma / std).reshape(-1, 1, 1, 1)), bn.bias - bn.running_mean * gamma / std 10 | 11 | 12 | def transII_addbranch(kernels, biases): 13 | return sum(kernels), sum(biases) 14 | 15 | 16 | def transIII_1x1_kxk(k1, b1, k2, b2, groups): 17 | if groups == 1: 18 | k = F.conv2d(k2, k1.permute(1, 0, 2, 3)) # 19 | b_hat = (k2 * b1.reshape(1, -1, 1, 1)).sum((1, 2, 3)) 20 | else: 21 | k_slices = [] 22 | b_slices = [] 23 | k1_T = k1.permute(1, 0, 2, 3) 24 | k1_group_width = k1.size(0) // groups 25 | k2_group_width = k2.size(0) // groups 26 | for g in range(groups): 27 | k1_T_slice = k1_T[:, g*k1_group_width:(g+1)*k1_group_width, :, :] 28 | k2_slice = k2[g*k2_group_width:(g+1)*k2_group_width, :, :, :] 29 | k_slices.append(F.conv2d(k2_slice, k1_T_slice)) 30 | b_slices.append((k2_slice * b1[g * k1_group_width:(g+1) * k1_group_width].reshape(1, -1, 1, 1)).sum((1, 2, 3))) 31 | k, b_hat = transIV_depthconcat(k_slices, b_slices) 32 | return k, b_hat + b2 33 | 34 | 35 | def transIV_depthconcat(kernels, biases): 36 | return torch.cat(kernels, dim=0), torch.cat(biases) 37 | 38 | 39 | def transV_avg(channels, kernel_size, groups): 40 | input_dim = channels // groups 41 | k = torch.zeros((channels, input_dim, kernel_size, kernel_size)) 42 | k[np.arange(channels), np.tile(np.arange(input_dim), groups), :, :] = 1.0 / kernel_size ** 2 43 | return k 44 | 45 | 46 | # This has not been tested with non-square kernels (kernel.size(2) != kernel.size(3)) nor even-size kernels 47 | def transVI_multiscale(kernel, target_kernel_size): 48 | H_pixels_to_pad = (target_kernel_size - kernel.size(2)) // 2 49 | W_pixels_to_pad = (target_kernel_size - kernel.size(3)) // 2 50 | return F.pad(kernel, [H_pixels_to_pad, H_pixels_to_pad, W_pixels_to_pad, W_pixels_to_pad]) 51 | -------------------------------------------------------------------------------- /yolov6/solver/build.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import os 4 | import math 5 | 6 | import torch 7 | import torch.nn as nn 8 | 9 | from yolov6.utils.events import LOGGER 10 | 11 | 12 | def build_optimizer(cfg, model): 13 | """ Build optimizer from cfg file.""" 14 | g_bnw, g_w, g_b = [], [], [] 15 | for v in model.modules(): 16 | if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): 17 | g_b.append(v.bias) 18 | if isinstance(v, nn.BatchNorm2d): 19 | g_bnw.append(v.weight) 20 | elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): 21 | g_w.append(v.weight) 22 | 23 | assert cfg.solver.optim == 'SGD' or 'Adam', 'ERROR: unknown optimizer, use SGD defaulted' 24 | if cfg.solver.optim == 'SGD': 25 | optimizer = torch.optim.SGD(g_bnw, lr=cfg.solver.lr0, momentum=cfg.solver.momentum, nesterov=True) 26 | elif cfg.solver.optim == 'Adam': 27 | optimizer = torch.optim.Adam(g_bnw, lr=cfg.solver.lr0, betas=(cfg.solver.momentum, 0.999)) 28 | 29 | optimizer.add_param_group({'params': g_w, 'weight_decay': cfg.solver.weight_decay}) 30 | optimizer.add_param_group({'params': g_b}) 31 | 32 | del g_bnw, g_w, g_b 33 | return optimizer 34 | 35 | 36 | def build_lr_scheduler(cfg, optimizer, epochs): 37 | """Build learning rate scheduler from cfg file.""" 38 | if cfg.solver.lr_scheduler == 'Cosine': 39 | lf = lambda x: ((1 - math.cos(x * math.pi / epochs)) / 2) * (cfg.solver.lrf - 1) + 1 40 | elif cfg.solver.lr_scheduler == 'Constant': 41 | lf = lambda x: 1.0 42 | else: 43 | LOGGER.error('unknown lr scheduler, use Cosine defaulted') 44 | 45 | scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) 46 | return scheduler, lf 47 | -------------------------------------------------------------------------------- /yolov6/utils/Arial.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/airockchip/YOLOv6/0e7c2d5a93f6d49ed5ab6f005ccdd9d9bbd3db9b/yolov6/utils/Arial.ttf -------------------------------------------------------------------------------- /yolov6/utils/checkpoint.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import os 4 | import shutil 5 | import torch 6 | import os.path as osp 7 | from yolov6.utils.events import LOGGER 8 | from yolov6.utils.torch_utils import fuse_model 9 | 10 | 11 | def load_state_dict(weights, model, map_location=None): 12 | """Load weights from checkpoint file, only assign weights those layers' name and shape are match.""" 13 | ckpt = torch.load(weights, map_location=map_location) 14 | state_dict = ckpt['model'].float().state_dict() 15 | model_state_dict = model.state_dict() 16 | state_dict = {k: v for k, v in state_dict.items() if k in model_state_dict and v.shape == model_state_dict[k].shape} 17 | model.load_state_dict(state_dict, strict=False) 18 | del ckpt, state_dict, model_state_dict 19 | return model 20 | 21 | 22 | def load_checkpoint(weights, map_location=None, inplace=True, fuse=True): 23 | """Load model from checkpoint file.""" 24 | LOGGER.info("Loading checkpoint from {}".format(weights)) 25 | ckpt = torch.load(weights, map_location=map_location) # load 26 | model = ckpt['ema' if ckpt.get('ema') else 'model'].float() 27 | if fuse: 28 | LOGGER.info("\nFusing model...") 29 | model = fuse_model(model).eval() 30 | else: 31 | model = model.eval() 32 | return model 33 | 34 | 35 | def save_checkpoint(ckpt, is_best, save_dir, model_name=""): 36 | """ Save checkpoint to the disk.""" 37 | if not osp.exists(save_dir): 38 | os.makedirs(save_dir) 39 | filename = osp.join(save_dir, model_name + '.pt') 40 | torch.save(ckpt, filename) 41 | if is_best: 42 | best_filename = osp.join(save_dir, 'best_ckpt.pt') 43 | shutil.copyfile(filename, best_filename) 44 | 45 | 46 | def strip_optimizer(ckpt_dir, epoch): 47 | """Delete optimizer from saved checkpoint file""" 48 | for s in ['best', 'last']: 49 | ckpt_path = osp.join(ckpt_dir, '{}_ckpt.pt'.format(s)) 50 | if not osp.exists(ckpt_path): 51 | continue 52 | ckpt = torch.load(ckpt_path, map_location=torch.device('cpu')) 53 | if ckpt.get('ema'): 54 | ckpt['model'] = ckpt['ema'] # replace model with ema 55 | for k in ['optimizer', 'ema', 'updates']: # keys 56 | ckpt[k] = None 57 | ckpt['epoch'] = epoch 58 | ckpt['model'].half() # to FP16 59 | for p in ckpt['model'].parameters(): 60 | p.requires_grad = False 61 | torch.save(ckpt, ckpt_path) 62 | -------------------------------------------------------------------------------- /yolov6/utils/config.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | # The code is based on 4 | # https://github.com/open-mmlab/mmcv/blob/master/mmcv/utils/config.py 5 | # Copyright (c) OpenMMLab. 6 | 7 | import os.path as osp 8 | import shutil 9 | import sys 10 | import tempfile 11 | from importlib import import_module 12 | from addict import Dict 13 | 14 | 15 | class ConfigDict(Dict): 16 | 17 | def __missing__(self, name): 18 | raise KeyError(name) 19 | 20 | def __getattr__(self, name): 21 | try: 22 | value = super(ConfigDict, self).__getattr__(name) 23 | except KeyError: 24 | ex = AttributeError("'{}' object has no attribute '{}'".format( 25 | self.__class__.__name__, name)) 26 | except Exception as e: 27 | ex = e 28 | else: 29 | return value 30 | raise ex 31 | 32 | 33 | class Config(object): 34 | 35 | @staticmethod 36 | def _file2dict(filename): 37 | filename = str(filename) 38 | if filename.endswith('.py'): 39 | with tempfile.TemporaryDirectory() as temp_config_dir: 40 | shutil.copyfile(filename, 41 | osp.join(temp_config_dir, '_tempconfig.py')) 42 | sys.path.insert(0, temp_config_dir) 43 | mod = import_module('_tempconfig') 44 | sys.path.pop(0) 45 | cfg_dict = { 46 | name: value 47 | for name, value in mod.__dict__.items() 48 | if not name.startswith('__') 49 | } 50 | # delete imported module 51 | del sys.modules['_tempconfig'] 52 | else: 53 | raise IOError('Only .py type are supported now!') 54 | cfg_text = filename + '\n' 55 | with open(filename, 'r') as f: 56 | cfg_text += f.read() 57 | 58 | return cfg_dict, cfg_text 59 | 60 | @staticmethod 61 | def fromfile(filename): 62 | cfg_dict, cfg_text = Config._file2dict(filename) 63 | return Config(cfg_dict, cfg_text=cfg_text, filename=filename) 64 | 65 | def __init__(self, cfg_dict=None, cfg_text=None, filename=None): 66 | if cfg_dict is None: 67 | cfg_dict = dict() 68 | elif not isinstance(cfg_dict, dict): 69 | raise TypeError('cfg_dict must be a dict, but got {}'.format( 70 | type(cfg_dict))) 71 | 72 | super(Config, self).__setattr__('_cfg_dict', ConfigDict(cfg_dict)) 73 | super(Config, self).__setattr__('_filename', filename) 74 | if cfg_text: 75 | text = cfg_text 76 | elif filename: 77 | with open(filename, 'r') as f: 78 | text = f.read() 79 | else: 80 | text = '' 81 | super(Config, self).__setattr__('_text', text) 82 | 83 | @property 84 | def filename(self): 85 | return self._filename 86 | 87 | @property 88 | def text(self): 89 | return self._text 90 | 91 | def __repr__(self): 92 | return 'Config (path: {}): {}'.format(self.filename, 93 | self._cfg_dict.__repr__()) 94 | 95 | def __getattr__(self, name): 96 | return getattr(self._cfg_dict, name) 97 | 98 | def __setattr__(self, name, value): 99 | if isinstance(value, dict): 100 | value = ConfigDict(value) 101 | self._cfg_dict.__setattr__(name, value) 102 | -------------------------------------------------------------------------------- /yolov6/utils/ema.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # The code is based on 4 | # https://github.com/ultralytics/yolov5/blob/master/utils/torch_utils.py 5 | import math 6 | from copy import deepcopy 7 | import torch 8 | import torch.nn as nn 9 | 10 | 11 | class ModelEMA: 12 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 13 | Keep a moving average of everything in the model state_dict (parameters and buffers). 14 | This is intended to allow functionality like 15 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 16 | A smoothed version of the weights is necessary for some training schemes to perform well. 17 | This class is sensitive where it is initialized in the sequence of model init, 18 | GPU assignment and distributed training wrappers. 19 | """ 20 | 21 | def __init__(self, model, decay=0.9999, updates=0): 22 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA 23 | self.updates = updates 24 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) 25 | for param in self.ema.parameters(): 26 | param.requires_grad_(False) 27 | 28 | def update(self, model): 29 | with torch.no_grad(): 30 | self.updates += 1 31 | decay = self.decay(self.updates) 32 | 33 | state_dict = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict 34 | for k, item in self.ema.state_dict().items(): 35 | if item.dtype.is_floating_point: 36 | item *= decay 37 | item += (1 - decay) * state_dict[k].detach() 38 | 39 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): 40 | copy_attr(self.ema, model, include, exclude) 41 | 42 | 43 | def copy_attr(a, b, include=(), exclude=()): 44 | """Copy attributes from one instance and set them to another instance.""" 45 | for k, item in b.__dict__.items(): 46 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 47 | continue 48 | else: 49 | setattr(a, k, item) 50 | 51 | 52 | def is_parallel(model): 53 | '''Return True if model's type is DP or DDP, else False.''' 54 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 55 | 56 | 57 | def de_parallel(model): 58 | '''De-parallelize a model. Return single-GPU model if model's type is DP or DDP.''' 59 | return model.module if is_parallel(model) else model 60 | -------------------------------------------------------------------------------- /yolov6/utils/envs.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import os 4 | import random 5 | import numpy as np 6 | 7 | import torch 8 | import torch.backends.cudnn as cudnn 9 | from yolov6.utils.events import LOGGER 10 | 11 | 12 | def get_envs(): 13 | """Get PyTorch needed environments from system envirionments.""" 14 | local_rank = int(os.getenv('LOCAL_RANK', -1)) 15 | rank = int(os.getenv('RANK', -1)) 16 | world_size = int(os.getenv('WORLD_SIZE', 1)) 17 | return local_rank, rank, world_size 18 | 19 | 20 | def select_device(device): 21 | """Set devices' information to the program. 22 | Args: 23 | device: a string, like 'cpu' or '1,2,3,4' 24 | Returns: 25 | torch.device 26 | """ 27 | if device == 'cpu': 28 | os.environ['CUDA_VISIBLE_DEVICES'] = '-1' 29 | LOGGER.info('Using CPU for training... ') 30 | elif device: 31 | os.environ['CUDA_VISIBLE_DEVICES'] = device 32 | assert torch.cuda.is_available() 33 | nd = len(device.strip().split(',')) 34 | LOGGER.info(f'Using {nd} GPU for training... ') 35 | cuda = device != 'cpu' and torch.cuda.is_available() 36 | device = torch.device('cuda:0' if cuda else 'cpu') 37 | return device 38 | 39 | 40 | def set_random_seed(seed, deterministic=False): 41 | """ Set random state to random libray, numpy, torch and cudnn. 42 | Args: 43 | seed: int value. 44 | deterministic: bool value. 45 | """ 46 | random.seed(seed) 47 | np.random.seed(seed) 48 | torch.manual_seed(seed) 49 | if deterministic: 50 | cudnn.deterministic = True 51 | cudnn.benchmark = False 52 | else: 53 | cudnn.deterministic = False 54 | cudnn.benchmark = True 55 | -------------------------------------------------------------------------------- /yolov6/utils/events.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import os 4 | import yaml 5 | import logging 6 | import shutil 7 | 8 | 9 | def set_logging(name=None): 10 | rank = int(os.getenv('RANK', -1)) 11 | logging.basicConfig(format="%(message)s", level=logging.INFO if (rank in (-1, 0)) else logging.WARNING) 12 | return logging.getLogger(name) 13 | 14 | 15 | LOGGER = set_logging(__name__) 16 | NCOLS = min(100, shutil.get_terminal_size().columns) 17 | 18 | 19 | def load_yaml(file_path): 20 | """Load data from yaml file.""" 21 | if isinstance(file_path, str): 22 | with open(file_path, errors='ignore') as f: 23 | data_dict = yaml.safe_load(f) 24 | return data_dict 25 | 26 | 27 | def save_yaml(data_dict, save_path): 28 | """Save data to yaml file""" 29 | with open(save_path, 'w') as f: 30 | yaml.safe_dump(data_dict, f, sort_keys=False) 31 | 32 | 33 | def write_tblog(tblogger, epoch, results, losses): 34 | """Display mAP and loss information to log.""" 35 | tblogger.add_scalar("val/mAP@0.5", results[0], epoch + 1) 36 | tblogger.add_scalar("val/mAP@0.50:0.95", results[1], epoch + 1) 37 | 38 | tblogger.add_scalar("train/iou_loss", losses[0], epoch + 1) 39 | tblogger.add_scalar("train/dist_focalloss", losses[1], epoch + 1) 40 | tblogger.add_scalar("train/cls_loss", losses[2], epoch + 1) 41 | 42 | tblogger.add_scalar("x/lr0", results[2], epoch + 1) 43 | tblogger.add_scalar("x/lr1", results[3], epoch + 1) 44 | tblogger.add_scalar("x/lr2", results[4], epoch + 1) 45 | 46 | 47 | def write_tbimg(tblogger, imgs, step, type='train'): 48 | """Display train_batch and validation predictions to tensorboard.""" 49 | if type == 'train': 50 | tblogger.add_image(f'train_batch', imgs, step + 1, dataformats='HWC') 51 | elif type == 'val': 52 | for idx, img in enumerate(imgs): 53 | tblogger.add_image(f'val_img_{idx + 1}', img, step + 1, dataformats='HWC') 54 | else: 55 | LOGGER.warning('WARNING: Unknown image type to visualize.\n') 56 | -------------------------------------------------------------------------------- /yolov6/utils/general.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import os 4 | import glob 5 | import torch 6 | import requests 7 | from pathlib import Path 8 | from yolov6.utils.events import LOGGER 9 | 10 | def increment_name(path): 11 | '''increase save directory's id''' 12 | path = Path(path) 13 | sep = '' 14 | if path.exists(): 15 | path, suffix = (path.with_suffix(''), path.suffix) if path.is_file() else (path, '') 16 | for n in range(1, 9999): 17 | p = f'{path}{sep}{n}{suffix}' 18 | if not os.path.exists(p): 19 | break 20 | path = Path(p) 21 | return path 22 | 23 | 24 | def find_latest_checkpoint(search_dir='.'): 25 | '''Find the most recent saved checkpoint in search_dir.''' 26 | checkpoint_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True) 27 | return max(checkpoint_list, key=os.path.getctime) if checkpoint_list else '' 28 | 29 | 30 | def dist2bbox(distance, anchor_points, box_format='xyxy'): 31 | '''Transform distance(ltrb) to box(xywh or xyxy).''' 32 | lt, rb = torch.split(distance, 2, -1) 33 | x1y1 = anchor_points - lt 34 | x2y2 = anchor_points + rb 35 | if box_format == 'xyxy': 36 | bbox = torch.cat([x1y1, x2y2], -1) 37 | elif box_format == 'xywh': 38 | c_xy = (x1y1 + x2y2) / 2 39 | wh = x2y2 - x1y1 40 | bbox = torch.cat([c_xy, wh], -1) 41 | return bbox 42 | 43 | 44 | def bbox2dist(anchor_points, bbox, reg_max): 45 | '''Transform bbox(xyxy) to dist(ltrb).''' 46 | x1y1, x2y2 = torch.split(bbox, 2, -1) 47 | lt = anchor_points - x1y1 48 | rb = x2y2 - anchor_points 49 | dist = torch.cat([lt, rb], -1).clip(0, reg_max - 0.01) 50 | return dist 51 | 52 | 53 | def xywh2xyxy(bboxes): 54 | '''Transform bbox(xywh) to box(xyxy).''' 55 | bboxes[..., 0] = bboxes[..., 0] - bboxes[..., 2] * 0.5 56 | bboxes[..., 1] = bboxes[..., 1] - bboxes[..., 3] * 0.5 57 | bboxes[..., 2] = bboxes[..., 0] + bboxes[..., 2] 58 | bboxes[..., 3] = bboxes[..., 1] + bboxes[..., 3] 59 | return bboxes 60 | 61 | 62 | def box_iou(box1, box2): 63 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 64 | """ 65 | Return intersection-over-union (Jaccard index) of boxes. 66 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 67 | Arguments: 68 | box1 (Tensor[N, 4]) 69 | box2 (Tensor[M, 4]) 70 | Returns: 71 | iou (Tensor[N, M]): the NxM matrix containing the pairwise 72 | IoU values for every element in boxes1 and boxes2 73 | """ 74 | 75 | def box_area(box): 76 | # box = 4xn 77 | return (box[2] - box[0]) * (box[3] - box[1]) 78 | 79 | area1 = box_area(box1.T) 80 | area2 = box_area(box2.T) 81 | 82 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 83 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 84 | return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) 85 | 86 | 87 | def download_ckpt(path): 88 | """Download checkpoints of the pretrained models""" 89 | basename = os.path.basename(path) 90 | dir = os.path.abspath(os.path.dirname(path)) 91 | os.makedirs(dir, exist_ok=True) 92 | LOGGER.info(f"checkpoint {basename} not exist, try to downloaded it from github.") 93 | # need to update the link with every release 94 | url = f"https://github.com/meituan/YOLOv6/releases/download/0.3.0/{basename}" 95 | r = requests.get(url, allow_redirects=True) 96 | assert r.status_code == 200, "Unable to download checkpoints, manually download it" 97 | open(path, 'wb').write(r.content) 98 | LOGGER.info(f"checkpoint {basename} downloaded and saved") 99 | -------------------------------------------------------------------------------- /yolov6/utils/nms.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # The code is based on 4 | # https://github.com/ultralytics/yolov5/blob/master/utils/general.py 5 | 6 | import os 7 | import time 8 | import numpy as np 9 | import cv2 10 | import torch 11 | import torchvision 12 | 13 | 14 | # Settings 15 | torch.set_printoptions(linewidth=320, precision=5, profile='long') 16 | np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 17 | cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) 18 | os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads 19 | 20 | 21 | def xywh2xyxy(x): 22 | '''Convert boxes with shape [n, 4] from [x, y, w, h] to [x1, y1, x2, y2] where x1y1 is top-left, x2y2=bottom-right.''' 23 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 24 | y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x 25 | y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y 26 | y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x 27 | y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y 28 | return y 29 | 30 | 31 | def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, max_det=300): 32 | """Runs Non-Maximum Suppression (NMS) on inference results. 33 | This code is borrowed from: https://github.com/ultralytics/yolov5/blob/47233e1698b89fc437a4fb9463c815e9171be955/utils/general.py#L775 34 | Args: 35 | prediction: (tensor), with shape [N, 5 + num_classes], N is the number of bboxes. 36 | conf_thres: (float) confidence threshold. 37 | iou_thres: (float) iou threshold. 38 | classes: (None or list[int]), if a list is provided, nms only keep the classes you provide. 39 | agnostic: (bool), when it is set to True, we do class-independent nms, otherwise, different class would do nms respectively. 40 | multi_label: (bool), when it is set to True, one box can have multi labels, otherwise, one box only huave one label. 41 | max_det:(int), max number of output bboxes. 42 | 43 | Returns: 44 | list of detections, echo item is one tensor with shape (num_boxes, 6), 6 is for [xyxy, conf, cls]. 45 | """ 46 | 47 | num_classes = prediction.shape[2] - 5 # number of classes 48 | pred_candidates = torch.logical_and(prediction[..., 4] > conf_thres, torch.max(prediction[..., 5:], axis=-1)[0] > conf_thres) # candidates 49 | # Check the parameters. 50 | assert 0 <= conf_thres <= 1, f'conf_thresh must be in 0.0 to 1.0, however {conf_thres} is provided.' 51 | assert 0 <= iou_thres <= 1, f'iou_thres must be in 0.0 to 1.0, however {iou_thres} is provided.' 52 | 53 | # Function settings. 54 | max_wh = 4096 # maximum box width and height 55 | max_nms = 30000 # maximum number of boxes put into torchvision.ops.nms() 56 | time_limit = 10.0 # quit the function when nms cost time exceed the limit time. 57 | multi_label &= num_classes > 1 # multiple labels per box 58 | 59 | tik = time.time() 60 | output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0] 61 | for img_idx, x in enumerate(prediction): # image index, image inference 62 | x = x[pred_candidates[img_idx]] # confidence 63 | 64 | # If no box remains, skip the next process. 65 | if not x.shape[0]: 66 | continue 67 | 68 | # confidence multiply the objectness 69 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 70 | 71 | # (center x, center y, width, height) to (x1, y1, x2, y2) 72 | box = xywh2xyxy(x[:, :4]) 73 | 74 | # Detections matrix's shape is (n,6), each row represents (xyxy, conf, cls) 75 | if multi_label: 76 | box_idx, class_idx = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 77 | x = torch.cat((box[box_idx], x[box_idx, class_idx + 5, None], class_idx[:, None].float()), 1) 78 | else: # Only keep the class with highest scores. 79 | conf, class_idx = x[:, 5:].max(1, keepdim=True) 80 | x = torch.cat((box, conf, class_idx.float()), 1)[conf.view(-1) > conf_thres] 81 | 82 | # Filter by class, only keep boxes whose category is in classes. 83 | if classes is not None: 84 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 85 | 86 | # Check shape 87 | num_box = x.shape[0] # number of boxes 88 | if not num_box: # no boxes kept. 89 | continue 90 | elif num_box > max_nms: # excess max boxes' number. 91 | x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence 92 | 93 | # Batched NMS 94 | class_offset = x[:, 5:6] * (0 if agnostic else max_wh) # classes 95 | boxes, scores = x[:, :4] + class_offset, x[:, 4] # boxes (offset by class), scores 96 | keep_box_idx = torchvision.ops.nms(boxes, scores, iou_thres) # NMS 97 | if keep_box_idx.shape[0] > max_det: # limit detections 98 | keep_box_idx = keep_box_idx[:max_det] 99 | 100 | output[img_idx] = x[keep_box_idx] 101 | if (time.time() - tik) > time_limit: 102 | print(f'WARNING: NMS cost time exceed the limited {time_limit}s.') 103 | break # time limit exceeded 104 | 105 | return output 106 | -------------------------------------------------------------------------------- /yolov6/utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | 4 | import time 5 | from contextlib import contextmanager 6 | from copy import deepcopy 7 | import torch 8 | import torch.distributed as dist 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | from yolov6.utils.events import LOGGER 12 | 13 | try: 14 | import thop # for FLOPs computation 15 | except ImportError: 16 | thop = None 17 | 18 | 19 | @contextmanager 20 | def torch_distributed_zero_first(local_rank: int): 21 | """ 22 | Decorator to make all processes in distributed training wait for each local_master to do something. 23 | """ 24 | if local_rank not in [-1, 0]: 25 | dist.barrier(device_ids=[local_rank]) 26 | yield 27 | if local_rank == 0: 28 | dist.barrier(device_ids=[0]) 29 | 30 | 31 | def time_sync(): 32 | '''Waits for all kernels in all streams on a CUDA device to complete if cuda is available.''' 33 | if torch.cuda.is_available(): 34 | torch.cuda.synchronize() 35 | return time.time() 36 | 37 | 38 | def initialize_weights(model): 39 | for m in model.modules(): 40 | t = type(m) 41 | if t is nn.Conv2d: 42 | pass 43 | elif t is nn.BatchNorm2d: 44 | m.eps = 1e-3 45 | m.momentum = 0.03 46 | elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]: 47 | m.inplace = True 48 | 49 | 50 | def fuse_conv_and_bn(conv, bn): 51 | '''Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/.''' 52 | fusedconv = ( 53 | nn.Conv2d( 54 | conv.in_channels, 55 | conv.out_channels, 56 | kernel_size=conv.kernel_size, 57 | stride=conv.stride, 58 | padding=conv.padding, 59 | groups=conv.groups, 60 | bias=True, 61 | ) 62 | .requires_grad_(False) 63 | .to(conv.weight.device) 64 | ) 65 | 66 | # prepare filters 67 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 68 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 69 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape)) 70 | 71 | # prepare spatial bias 72 | b_conv = ( 73 | torch.zeros(conv.weight.size(0), device=conv.weight.device) 74 | if conv.bias is None 75 | else conv.bias 76 | ) 77 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div( 78 | torch.sqrt(bn.running_var + bn.eps) 79 | ) 80 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) 81 | 82 | return fusedconv 83 | 84 | 85 | def fuse_model(model): 86 | '''Fuse convolution and batchnorm layers of the model.''' 87 | from yolov6.layers.common import Conv, SimConv, Conv_C3 88 | 89 | for m in model.modules(): 90 | if (type(m) is Conv or type(m) is SimConv or type(m) is Conv_C3) and hasattr(m, "bn"): 91 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 92 | delattr(m, "bn") # remove batchnorm 93 | m.forward = m.forward_fuse # update forward 94 | return model 95 | 96 | 97 | def get_model_info(model, img_size=640): 98 | """Get model Params and GFlops. 99 | Code base on https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/utils/model_utils.py 100 | """ 101 | from thop import profile 102 | stride = 64 #32 103 | img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device) 104 | 105 | flops, params = profile(deepcopy(model), inputs=(img,), verbose=False) 106 | params /= 1e6 107 | flops /= 1e9 108 | img_size = img_size if isinstance(img_size, list) else [img_size, img_size] 109 | flops *= img_size[0] * img_size[1] / stride / stride * 2 # Gflops 110 | info = "Params: {:.2f}M, Gflops: {:.2f}".format(params, flops) 111 | return info 112 | --------------------------------------------------------------------------------