├── .dockerignore ├── .gitattributes ├── .github ├── ISSUE_TEMPLATE │ ├── --bug-report.md │ ├── --feature-request.md │ └── -question.md └── workflows │ ├── ci-testing.yml │ ├── greetings.yml │ ├── rebase.yml │ └── stale.yml ├── .gitignore ├── CIoU.png ├── Dockerfile ├── LICENSE ├── README.md ├── data ├── coco.yaml ├── coco128.yaml ├── hyp.finetune.yaml ├── hyp.scratch.yaml ├── scripts │ ├── get_coco.sh │ └── get_voc.sh └── voc.yaml ├── detect.py ├── hubconf.py ├── inference └── images │ ├── bus.jpg │ └── zidane.jpg ├── models ├── __init__.py ├── common.py ├── experimental.py ├── export.py ├── hub │ ├── yolov3-spp.yaml │ ├── yolov5-fpn.yaml │ └── yolov5-panet.yaml ├── yolo.py ├── yolov5l.yaml ├── yolov5m.yaml ├── yolov5s.yaml └── yolov5x.yaml ├── requirements.txt ├── sotabench.py ├── test.py ├── train.py ├── tutorial.ipynb ├── utils ├── __init__.py ├── activations.py ├── datasets.py ├── evolve.sh ├── general.py ├── google_utils.py └── torch_utils.py └── weights └── download_weights.sh /.dockerignore: -------------------------------------------------------------------------------- 1 | # Repo-specific DockerIgnore ------------------------------------------------------------------------------------------- 2 | .git 3 | .cache 4 | .idea 5 | runs 6 | output 7 | coco 8 | storage.googleapis.com 9 | 10 | data/samples/* 11 | **/results*.txt 12 | *.jpg 13 | 14 | # Neural Network weights ----------------------------------------------------------------------------------------------- 15 | **/*.weights 16 | **/*.pt 17 | **/*.pth 18 | **/*.onnx 19 | **/*.mlmodel 20 | **/*.torchscript 21 | 22 | 23 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 24 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 25 | 26 | 27 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 28 | # Byte-compiled / optimized / DLL files 29 | __pycache__/ 30 | *.py[cod] 31 | *$py.class 32 | 33 | # C extensions 34 | *.so 35 | 36 | # Distribution / packaging 37 | .Python 38 | env/ 39 | build/ 40 | develop-eggs/ 41 | dist/ 42 | downloads/ 43 | eggs/ 44 | .eggs/ 45 | lib/ 46 | lib64/ 47 | parts/ 48 | sdist/ 49 | var/ 50 | wheels/ 51 | *.egg-info/ 52 | .installed.cfg 53 | *.egg 54 | 55 | # PyInstaller 56 | # Usually these files are written by a python script from a template 57 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 58 | *.manifest 59 | *.spec 60 | 61 | # Installer logs 62 | pip-log.txt 63 | pip-delete-this-directory.txt 64 | 65 | # Unit test / coverage reports 66 | htmlcov/ 67 | .tox/ 68 | .coverage 69 | .coverage.* 70 | .cache 71 | nosetests.xml 72 | coverage.xml 73 | *.cover 74 | .hypothesis/ 75 | 76 | # Translations 77 | *.mo 78 | *.pot 79 | 80 | # Django stuff: 81 | *.log 82 | local_settings.py 83 | 84 | # Flask stuff: 85 | instance/ 86 | .webassets-cache 87 | 88 | # Scrapy stuff: 89 | .scrapy 90 | 91 | # Sphinx documentation 92 | docs/_build/ 93 | 94 | # PyBuilder 95 | target/ 96 | 97 | # Jupyter Notebook 98 | .ipynb_checkpoints 99 | 100 | # pyenv 101 | .python-version 102 | 103 | # celery beat schedule file 104 | celerybeat-schedule 105 | 106 | # SageMath parsed files 107 | *.sage.py 108 | 109 | # dotenv 110 | .env 111 | 112 | # virtualenv 113 | .venv 114 | venv*/ 115 | ENV/ 116 | 117 | # Spyder project settings 118 | .spyderproject 119 | .spyproject 120 | 121 | # Rope project settings 122 | .ropeproject 123 | 124 | # mkdocs documentation 125 | /site 126 | 127 | # mypy 128 | .mypy_cache/ 129 | 130 | 131 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 132 | 133 | # General 134 | .DS_Store 135 | .AppleDouble 136 | .LSOverride 137 | 138 | # Icon must end with two \r 139 | Icon 140 | Icon? 141 | 142 | # Thumbnails 143 | ._* 144 | 145 | # Files that might appear in the root of a volume 146 | .DocumentRevisions-V100 147 | .fseventsd 148 | .Spotlight-V100 149 | .TemporaryItems 150 | .Trashes 151 | .VolumeIcon.icns 152 | .com.apple.timemachine.donotpresent 153 | 154 | # Directories potentially created on remote AFP share 155 | .AppleDB 156 | .AppleDesktop 157 | Network Trash Folder 158 | Temporary Items 159 | .apdisk 160 | 161 | 162 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 163 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 164 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 165 | 166 | # User-specific stuff: 167 | .idea/* 168 | .idea/**/workspace.xml 169 | .idea/**/tasks.xml 170 | .idea/dictionaries 171 | .html # Bokeh Plots 172 | .pg # TensorFlow Frozen Graphs 173 | .avi # videos 174 | 175 | # Sensitive or high-churn files: 176 | .idea/**/dataSources/ 177 | .idea/**/dataSources.ids 178 | .idea/**/dataSources.local.xml 179 | .idea/**/sqlDataSources.xml 180 | .idea/**/dynamic.xml 181 | .idea/**/uiDesigner.xml 182 | 183 | # Gradle: 184 | .idea/**/gradle.xml 185 | .idea/**/libraries 186 | 187 | # CMake 188 | cmake-build-debug/ 189 | cmake-build-release/ 190 | 191 | # Mongo Explorer plugin: 192 | .idea/**/mongoSettings.xml 193 | 194 | ## File-based project format: 195 | *.iws 196 | 197 | ## Plugin-specific files: 198 | 199 | # IntelliJ 200 | out/ 201 | 202 | # mpeltonen/sbt-idea plugin 203 | .idea_modules/ 204 | 205 | # JIRA plugin 206 | atlassian-ide-plugin.xml 207 | 208 | # Cursive Clojure plugin 209 | .idea/replstate.xml 210 | 211 | # Crashlytics plugin (for Android Studio and IntelliJ) 212 | com_crashlytics_export_strings.xml 213 | crashlytics.properties 214 | crashlytics-build.properties 215 | fabric.properties 216 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # this drop notebooks from GitHub language stats 2 | *.ipynb linguist-vendored 3 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/--bug-report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "\U0001F41BBug report" 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: bug 6 | assignees: '' 7 | 8 | --- 9 | 10 | Before submitting a bug report, please be aware that your issue **must be reproducible** with all of the following, otherwise it is non-actionable, and we can not help you: 11 | - **Current repo**: run `git fetch && git status -uno` to check and `git pull` to update repo 12 | - **Common dataset**: coco.yaml or coco128.yaml 13 | - **Common environment**: Colab, Google Cloud, or Docker image. See https://github.com/ultralytics/yolov5#environments 14 | 15 | If this is a custom dataset/training question you **must include** your `train*.jpg`, `test*.jpg` and `results.png` figures, or we can not help you. You can generate these with `utils.plot_results()`. 16 | 17 | 18 | ## 🐛 Bug 19 | A clear and concise description of what the bug is. 20 | 21 | 22 | ## To Reproduce (REQUIRED) 23 | 24 | Input: 25 | ``` 26 | import torch 27 | 28 | a = torch.tensor([5]) 29 | c = a / 0 30 | ``` 31 | 32 | Output: 33 | ``` 34 | Traceback (most recent call last): 35 | File "/Users/glennjocher/opt/anaconda3/envs/env1/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code 36 | exec(code_obj, self.user_global_ns, self.user_ns) 37 | File "", line 5, in 38 | c = a / 0 39 | RuntimeError: ZeroDivisionError 40 | ``` 41 | 42 | 43 | ## Expected behavior 44 | A clear and concise description of what you expected to happen. 45 | 46 | 47 | ## Environment 48 | If applicable, add screenshots to help explain your problem. 49 | 50 | - OS: [e.g. Ubuntu] 51 | - GPU [e.g. 2080 Ti] 52 | 53 | 54 | ## Additional context 55 | Add any other context about the problem here. 56 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/--feature-request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "\U0001F680Feature request" 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: enhancement 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## 🚀 Feature 11 | 12 | 13 | ## Motivation 14 | 15 | 16 | 17 | ## Pitch 18 | 19 | 20 | 21 | ## Alternatives 22 | 23 | 24 | 25 | ## Additional context 26 | 27 | 28 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/-question.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "❓Question" 3 | about: Ask a general question 4 | title: '' 5 | labels: question 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## ❔Question 11 | 12 | 13 | ## Additional context 14 | -------------------------------------------------------------------------------- /.github/workflows/ci-testing.yml: -------------------------------------------------------------------------------- 1 | name: CI CPU testing 2 | 3 | on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows 4 | push: 5 | pull_request: 6 | schedule: 7 | - cron: "0 0 * * *" 8 | 9 | jobs: 10 | cpu-tests: 11 | 12 | runs-on: ${{ matrix.os }} 13 | strategy: 14 | fail-fast: false 15 | matrix: 16 | os: [ubuntu-latest, macos-latest, windows-latest] 17 | python-version: [3.8] 18 | model: ['yolov5s'] # models to test 19 | 20 | # Timeout: https://stackoverflow.com/a/59076067/4521646 21 | timeout-minutes: 50 22 | steps: 23 | - uses: actions/checkout@v2 24 | - name: Set up Python ${{ matrix.python-version }} 25 | uses: actions/setup-python@v2 26 | with: 27 | python-version: ${{ matrix.python-version }} 28 | 29 | # Note: This uses an internal pip API and may not always work 30 | # https://github.com/actions/cache/blob/master/examples.md#multiple-oss-in-a-workflow 31 | - name: Get pip cache 32 | id: pip-cache 33 | run: | 34 | python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)" 35 | 36 | - name: Cache pip 37 | uses: actions/cache@v1 38 | with: 39 | path: ${{ steps.pip-cache.outputs.dir }} 40 | key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }} 41 | restore-keys: | 42 | ${{ runner.os }}-${{ matrix.python-version }}-pip- 43 | 44 | - name: Install dependencies 45 | run: | 46 | python -m pip install --upgrade pip 47 | pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html 48 | pip install -q onnx 49 | python --version 50 | pip --version 51 | pip list 52 | shell: bash 53 | 54 | - name: Download data 55 | run: | 56 | curl -L -o temp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 57 | unzip -q temp.zip -d ../ 58 | rm temp.zip 59 | 60 | - name: Tests workflow 61 | run: | 62 | export PYTHONPATH="$PWD" # to run *.py. files in subdirectories 63 | di=cpu # inference devices # define device 64 | 65 | # train 66 | python train.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di 67 | # detect 68 | python detect.py --weights weights/${{ matrix.model }}.pt --device $di 69 | python detect.py --weights runs/exp0/weights/last.pt --device $di 70 | # test 71 | python test.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --device $di 72 | python test.py --img 256 --batch 8 --weights runs/exp0/weights/last.pt --device $di 73 | 74 | python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect 75 | python models/export.py --img 256 --batch 1 --weights weights/${{ matrix.model }}.pt # export 76 | shell: bash 77 | -------------------------------------------------------------------------------- /.github/workflows/greetings.yml: -------------------------------------------------------------------------------- 1 | name: Greetings 2 | 3 | on: [pull_request_target, issues] 4 | 5 | jobs: 6 | greeting: 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/first-interaction@v1 10 | with: 11 | repo-token: ${{ secrets.GITHUB_TOKEN }} 12 | pr-message: | 13 | Hello @${{ github.actor }}, thank you for submitting a PR! To allow your work to be integrated as seamlessly as possible, we advise you to: 14 | - Verify your PR is **up-to-date with origin/master.** If your PR is behind origin/master update by running the following, replacing 'feature' with the name of your local branch: 15 | ```bash 16 | git remote add upstream https://github.com/ultralytics/yolov5.git 17 | git fetch upstream 18 | git checkout feature # <----- replace 'feature' with local branch name 19 | git rebase upstream/master 20 | git push -u origin -f 21 | ``` 22 | - Verify all Continuous Integration (CI) **checks are passing**. 23 | - Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ -Bruce Lee 24 | 25 | issue-message: | 26 | Hello @${{ github.actor }}, thank you for your interest in our work! Please visit our [Custom Training Tutorial](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data) to get started, and see our [Jupyter Notebook](https://github.com/ultralytics/yolov5/blob/master/tutorial.ipynb) Open In Colab, [Docker Image](https://hub.docker.com/r/ultralytics/yolov5), and [Google Cloud Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) for example environments. 27 | 28 | If this is a bug report, please provide screenshots and **minimum viable code to reproduce your issue**, otherwise we can not help you. 29 | 30 | If this is a custom model or data training question, please note Ultralytics does **not** provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as: 31 | - **Cloud-based AI** systems operating on **hundreds of HD video streams in realtime.** 32 | - **Edge AI** integrated into custom iOS and Android apps for realtime **30 FPS video inference.** 33 | - **Custom data training**, hyperparameter evolution, and model exportation to any destination. 34 | 35 | For more information please visit https://www.ultralytics.com. 36 | -------------------------------------------------------------------------------- /.github/workflows/rebase.yml: -------------------------------------------------------------------------------- 1 | name: Automatic Rebase 2 | # https://github.com/marketplace/actions/automatic-rebase 3 | 4 | on: 5 | issue_comment: 6 | types: [created] 7 | 8 | jobs: 9 | rebase: 10 | name: Rebase 11 | if: github.event.issue.pull_request != '' && contains(github.event.comment.body, '/rebase') 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout the latest code 15 | uses: actions/checkout@v2 16 | with: 17 | fetch-depth: 0 18 | - name: Automatic Rebase 19 | uses: cirrus-actions/rebase@1.3.1 20 | env: 21 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 22 | -------------------------------------------------------------------------------- /.github/workflows/stale.yml: -------------------------------------------------------------------------------- 1 | name: Close stale issues 2 | on: 3 | schedule: 4 | - cron: "0 0 * * *" 5 | 6 | jobs: 7 | stale: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/stale@v1 11 | with: 12 | repo-token: ${{ secrets.GITHUB_TOKEN }} 13 | stale-issue-message: 'This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.' 14 | stale-pr-message: 'This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.' 15 | days-before-stale: 30 16 | days-before-close: 5 17 | exempt-issue-label: 'documentation,tutorial' 18 | operations-per-run: 100 # The maximum number of operations per run, used to control rate limiting. 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Repo-specific GitIgnore ---------------------------------------------------------------------------------------------- 2 | *.jpg 3 | *.jpeg 4 | *.png 5 | *.bmp 6 | *.tif 7 | *.tiff 8 | *.heic 9 | *.JPG 10 | *.JPEG 11 | *.PNG 12 | *.BMP 13 | *.TIF 14 | *.TIFF 15 | *.HEIC 16 | *.mp4 17 | *.mov 18 | *.MOV 19 | *.avi 20 | *.data 21 | *.json 22 | 23 | *.cfg 24 | !cfg/yolov3*.cfg 25 | 26 | storage.googleapis.com 27 | runs/* 28 | data/* 29 | !data/samples/zidane.jpg 30 | !data/samples/bus.jpg 31 | !data/coco.names 32 | !data/coco_paper.names 33 | !data/coco.data 34 | !data/coco_*.data 35 | !data/coco_*.txt 36 | !data/trainvalno5k.shapes 37 | !data/*.sh 38 | 39 | pycocotools/* 40 | results*.txt 41 | gcp_test*.sh 42 | 43 | # MATLAB GitIgnore ----------------------------------------------------------------------------------------------------- 44 | *.m~ 45 | *.mat 46 | !targets*.mat 47 | 48 | # Neural Network weights ----------------------------------------------------------------------------------------------- 49 | *.weights 50 | *.pt 51 | *.onnx 52 | *.mlmodel 53 | *.torchscript 54 | darknet53.conv.74 55 | yolov3-tiny.conv.15 56 | 57 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 58 | # Byte-compiled / optimized / DLL files 59 | __pycache__/ 60 | *.py[cod] 61 | *$py.class 62 | 63 | # C extensions 64 | *.so 65 | 66 | # Distribution / packaging 67 | .Python 68 | env/ 69 | build/ 70 | develop-eggs/ 71 | dist/ 72 | downloads/ 73 | eggs/ 74 | .eggs/ 75 | lib/ 76 | lib64/ 77 | parts/ 78 | sdist/ 79 | var/ 80 | wheels/ 81 | *.egg-info/ 82 | .installed.cfg 83 | *.egg 84 | 85 | # PyInstaller 86 | # Usually these files are written by a python script from a template 87 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 88 | *.manifest 89 | *.spec 90 | 91 | # Installer logs 92 | pip-log.txt 93 | pip-delete-this-directory.txt 94 | 95 | # Unit test / coverage reports 96 | htmlcov/ 97 | .tox/ 98 | .coverage 99 | .coverage.* 100 | .cache 101 | nosetests.xml 102 | coverage.xml 103 | *.cover 104 | .hypothesis/ 105 | 106 | # Translations 107 | *.mo 108 | *.pot 109 | 110 | # Django stuff: 111 | *.log 112 | local_settings.py 113 | 114 | # Flask stuff: 115 | instance/ 116 | .webassets-cache 117 | 118 | # Scrapy stuff: 119 | .scrapy 120 | 121 | # Sphinx documentation 122 | docs/_build/ 123 | 124 | # PyBuilder 125 | target/ 126 | 127 | # Jupyter Notebook 128 | .ipynb_checkpoints 129 | 130 | # pyenv 131 | .python-version 132 | 133 | # celery beat schedule file 134 | celerybeat-schedule 135 | 136 | # SageMath parsed files 137 | *.sage.py 138 | 139 | # dotenv 140 | .env 141 | 142 | # virtualenv 143 | .venv 144 | venv/ 145 | ENV/ 146 | 147 | # Spyder project settings 148 | .spyderproject 149 | .spyproject 150 | 151 | # Rope project settings 152 | .ropeproject 153 | 154 | # mkdocs documentation 155 | /site 156 | 157 | # mypy 158 | .mypy_cache/ 159 | 160 | 161 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 162 | 163 | # General 164 | .DS_Store 165 | .AppleDouble 166 | .LSOverride 167 | 168 | # Icon must end with two \r 169 | Icon 170 | Icon? 171 | 172 | # Thumbnails 173 | ._* 174 | 175 | # Files that might appear in the root of a volume 176 | .DocumentRevisions-V100 177 | .fseventsd 178 | .Spotlight-V100 179 | .TemporaryItems 180 | .Trashes 181 | .VolumeIcon.icns 182 | .com.apple.timemachine.donotpresent 183 | 184 | # Directories potentially created on remote AFP share 185 | .AppleDB 186 | .AppleDesktop 187 | Network Trash Folder 188 | Temporary Items 189 | .apdisk 190 | 191 | 192 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 193 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 194 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 195 | 196 | # User-specific stuff: 197 | .idea/* 198 | .idea/**/workspace.xml 199 | .idea/**/tasks.xml 200 | .idea/dictionaries 201 | .html # Bokeh Plots 202 | .pg # TensorFlow Frozen Graphs 203 | .avi # videos 204 | 205 | # Sensitive or high-churn files: 206 | .idea/**/dataSources/ 207 | .idea/**/dataSources.ids 208 | .idea/**/dataSources.local.xml 209 | .idea/**/sqlDataSources.xml 210 | .idea/**/dynamic.xml 211 | .idea/**/uiDesigner.xml 212 | 213 | # Gradle: 214 | .idea/**/gradle.xml 215 | .idea/**/libraries 216 | 217 | # CMake 218 | cmake-build-debug/ 219 | cmake-build-release/ 220 | 221 | # Mongo Explorer plugin: 222 | .idea/**/mongoSettings.xml 223 | 224 | ## File-based project format: 225 | *.iws 226 | 227 | ## Plugin-specific files: 228 | 229 | # IntelliJ 230 | out/ 231 | 232 | # mpeltonen/sbt-idea plugin 233 | .idea_modules/ 234 | 235 | # JIRA plugin 236 | atlassian-ide-plugin.xml 237 | 238 | # Cursive Clojure plugin 239 | .idea/replstate.xml 240 | 241 | # Crashlytics plugin (for Android Studio and IntelliJ) 242 | com_crashlytics_export_strings.xml 243 | crashlytics.properties 244 | crashlytics-build.properties 245 | fabric.properties 246 | -------------------------------------------------------------------------------- /CIoU.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zzh-tju/yolov5/fca5b7b7a85bcadac394f99c1805e54d9c20a21f/CIoU.png -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch 2 | FROM nvcr.io/nvidia/pytorch:20.08-py3 3 | 4 | # Install dependencies 5 | RUN pip install --upgrade pip 6 | # COPY requirements.txt . 7 | # RUN pip install -r requirements.txt 8 | RUN pip install gsutil 9 | 10 | # Create working directory 11 | RUN mkdir -p /usr/src/app 12 | WORKDIR /usr/src/app 13 | 14 | # Copy contents 15 | COPY . /usr/src/app 16 | 17 | # Copy weights 18 | #RUN python3 -c "from models import *; \ 19 | #attempt_download('weights/yolov5s.pt'); \ 20 | #attempt_download('weights/yolov5m.pt'); \ 21 | #attempt_download('weights/yolov5l.pt')" 22 | 23 | 24 | # --------------------------------------------------- Extras Below --------------------------------------------------- 25 | 26 | # Build and Push 27 | # t=ultralytics/yolov5:latest && sudo docker build -t $t . && sudo docker push $t 28 | # for v in {300..303}; do t=ultralytics/coco:v$v && sudo docker build -t $t . && sudo docker push $t; done 29 | 30 | # Pull and Run 31 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host $t 32 | 33 | # Pull and Run with local directory access 34 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/coco:/usr/src/coco $t 35 | 36 | # Kill all 37 | # sudo docker kill "$(sudo docker ps -q)" 38 | 39 | # Kill all image-based 40 | # sudo docker kill $(sudo docker ps -a -q --filter ancestor=ultralytics/yolov5:latest) 41 | 42 | # Bash into running container 43 | # sudo docker container exec -it ba65811811ab bash 44 | 45 | # Bash into stopped container 46 | # sudo docker commit 092b16b25c5b usr/resume && sudo docker run -it --gpus all --ipc=host -v "$(pwd)"/coco:/usr/src/coco --entrypoint=sh usr/resume 47 | 48 | # Send weights to GCP 49 | # python -c "from utils.general import *; strip_optimizer('runs/exp0_*/weights/best.pt', 'tmp.pt')" && gsutil cp tmp.pt gs://*.pt 50 | 51 | # Clean up 52 | # docker system prune -a --volumes 53 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | GNU GENERAL PUBLIC LICENSE 2 | Version 3, 29 June 2007 3 | 4 | Copyright (C) 2007 Free Software Foundation, Inc. 5 | Everyone is permitted to copy and distribute verbatim copies 6 | of this license document, but changing it is not allowed. 7 | 8 | Preamble 9 | 10 | The GNU General Public License is a free, copyleft license for 11 | software and other kinds of works. 12 | 13 | The licenses for most software and other practical works are designed 14 | to take away your freedom to share and change the works. By contrast, 15 | the GNU General Public License is intended to guarantee your freedom to 16 | share and change all versions of a program--to make sure it remains free 17 | software for all its users. We, the Free Software Foundation, use the 18 | GNU General Public License for most of our software; it applies also to 19 | any other work released this way by its authors. You can apply it to 20 | your programs, too. 21 | 22 | When we speak of free software, we are referring to freedom, not 23 | price. Our General Public Licenses are designed to make sure that you 24 | have the freedom to distribute copies of free software (and charge for 25 | them if you wish), that you receive source code or can get it if you 26 | want it, that you can change the software or use pieces of it in new 27 | free programs, and that you know you can do these things. 28 | 29 | To protect your rights, we need to prevent others from denying you 30 | these rights or asking you to surrender the rights. Therefore, you have 31 | certain responsibilities if you distribute copies of the software, or if 32 | you modify it: responsibilities to respect the freedom of others. 33 | 34 | For example, if you distribute copies of such a program, whether 35 | gratis or for a fee, you must pass on to the recipients the same 36 | freedoms that you received. You must make sure that they, too, receive 37 | or can get the source code. And you must show them these terms so they 38 | know their rights. 39 | 40 | Developers that use the GNU GPL protect your rights with two steps: 41 | (1) assert copyright on the software, and (2) offer you this License 42 | giving you legal permission to copy, distribute and/or modify it. 43 | 44 | For the developers' and authors' protection, the GPL clearly explains 45 | that there is no warranty for this free software. For both users' and 46 | authors' sake, the GPL requires that modified versions be marked as 47 | changed, so that their problems will not be attributed erroneously to 48 | authors of previous versions. 49 | 50 | Some devices are designed to deny users access to install or run 51 | modified versions of the software inside them, although the manufacturer 52 | can do so. This is fundamentally incompatible with the aim of 53 | protecting users' freedom to change the software. The systematic 54 | pattern of such abuse occurs in the area of products for individuals to 55 | use, which is precisely where it is most unacceptable. Therefore, we 56 | have designed this version of the GPL to prohibit the practice for those 57 | products. If such problems arise substantially in other domains, we 58 | stand ready to extend this provision to those domains in future versions 59 | of the GPL, as needed to protect the freedom of users. 60 | 61 | Finally, every program is threatened constantly by software patents. 62 | States should not allow patents to restrict development and use of 63 | software on general-purpose computers, but in those that do, we wish to 64 | avoid the special danger that patents applied to a free program could 65 | make it effectively proprietary. To prevent this, the GPL assures that 66 | patents cannot be used to render the program non-free. 67 | 68 | The precise terms and conditions for copying, distribution and 69 | modification follow. 70 | 71 | TERMS AND CONDITIONS 72 | 73 | 0. Definitions. 74 | 75 | "This License" refers to version 3 of the GNU General Public License. 76 | 77 | "Copyright" also means copyright-like laws that apply to other kinds of 78 | works, such as semiconductor masks. 79 | 80 | "The Program" refers to any copyrightable work licensed under this 81 | License. Each licensee is addressed as "you". "Licensees" and 82 | "recipients" may be individuals or organizations. 83 | 84 | To "modify" a work means to copy from or adapt all or part of the work 85 | in a fashion requiring copyright permission, other than the making of an 86 | exact copy. The resulting work is called a "modified version" of the 87 | earlier work or a work "based on" the earlier work. 88 | 89 | A "covered work" means either the unmodified Program or a work based 90 | on the Program. 91 | 92 | To "propagate" a work means to do anything with it that, without 93 | permission, would make you directly or secondarily liable for 94 | infringement under applicable copyright law, except executing it on a 95 | computer or modifying a private copy. Propagation includes copying, 96 | distribution (with or without modification), making available to the 97 | public, and in some countries other activities as well. 98 | 99 | To "convey" a work means any kind of propagation that enables other 100 | parties to make or receive copies. Mere interaction with a user through 101 | a computer network, with no transfer of a copy, is not conveying. 102 | 103 | An interactive user interface displays "Appropriate Legal Notices" 104 | to the extent that it includes a convenient and prominently visible 105 | feature that (1) displays an appropriate copyright notice, and (2) 106 | tells the user that there is no warranty for the work (except to the 107 | extent that warranties are provided), that licensees may convey the 108 | work under this License, and how to view a copy of this License. If 109 | the interface presents a list of user commands or options, such as a 110 | menu, a prominent item in the list meets this criterion. 111 | 112 | 1. Source Code. 113 | 114 | The "source code" for a work means the preferred form of the work 115 | for making modifications to it. "Object code" means any non-source 116 | form of a work. 117 | 118 | A "Standard Interface" means an interface that either is an official 119 | standard defined by a recognized standards body, or, in the case of 120 | interfaces specified for a particular programming language, one that 121 | is widely used among developers working in that language. 122 | 123 | The "System Libraries" of an executable work include anything, other 124 | than the work as a whole, that (a) is included in the normal form of 125 | packaging a Major Component, but which is not part of that Major 126 | Component, and (b) serves only to enable use of the work with that 127 | Major Component, or to implement a Standard Interface for which an 128 | implementation is available to the public in source code form. A 129 | "Major Component", in this context, means a major essential component 130 | (kernel, window system, and so on) of the specific operating system 131 | (if any) on which the executable work runs, or a compiler used to 132 | produce the work, or an object code interpreter used to run it. 133 | 134 | The "Corresponding Source" for a work in object code form means all 135 | the source code needed to generate, install, and (for an executable 136 | work) run the object code and to modify the work, including scripts to 137 | control those activities. However, it does not include the work's 138 | System Libraries, or general-purpose tools or generally available free 139 | programs which are used unmodified in performing those activities but 140 | which are not part of the work. For example, Corresponding Source 141 | includes interface definition files associated with source files for 142 | the work, and the source code for shared libraries and dynamically 143 | linked subprograms that the work is specifically designed to require, 144 | such as by intimate data communication or control flow between those 145 | subprograms and other parts of the work. 146 | 147 | The Corresponding Source need not include anything that users 148 | can regenerate automatically from other parts of the Corresponding 149 | Source. 150 | 151 | The Corresponding Source for a work in source code form is that 152 | same work. 153 | 154 | 2. Basic Permissions. 155 | 156 | All rights granted under this License are granted for the term of 157 | copyright on the Program, and are irrevocable provided the stated 158 | conditions are met. This License explicitly affirms your unlimited 159 | permission to run the unmodified Program. The output from running a 160 | covered work is covered by this License only if the output, given its 161 | content, constitutes a covered work. This License acknowledges your 162 | rights of fair use or other equivalent, as provided by copyright law. 163 | 164 | You may make, run and propagate covered works that you do not 165 | convey, without conditions so long as your license otherwise remains 166 | in force. You may convey covered works to others for the sole purpose 167 | of having them make modifications exclusively for you, or provide you 168 | with facilities for running those works, provided that you comply with 169 | the terms of this License in conveying all material for which you do 170 | not control copyright. Those thus making or running the covered works 171 | for you must do so exclusively on your behalf, under your direction 172 | and control, on terms that prohibit them from making any copies of 173 | your copyrighted material outside their relationship with you. 174 | 175 | Conveying under any other circumstances is permitted solely under 176 | the conditions stated below. Sublicensing is not allowed; section 10 177 | makes it unnecessary. 178 | 179 | 3. Protecting Users' Legal Rights From Anti-Circumvention Law. 180 | 181 | No covered work shall be deemed part of an effective technological 182 | measure under any applicable law fulfilling obligations under article 183 | 11 of the WIPO copyright treaty adopted on 20 December 1996, or 184 | similar laws prohibiting or restricting circumvention of such 185 | measures. 186 | 187 | When you convey a covered work, you waive any legal power to forbid 188 | circumvention of technological measures to the extent such circumvention 189 | is effected by exercising rights under this License with respect to 190 | the covered work, and you disclaim any intention to limit operation or 191 | modification of the work as a means of enforcing, against the work's 192 | users, your or third parties' legal rights to forbid circumvention of 193 | technological measures. 194 | 195 | 4. Conveying Verbatim Copies. 196 | 197 | You may convey verbatim copies of the Program's source code as you 198 | receive it, in any medium, provided that you conspicuously and 199 | appropriately publish on each copy an appropriate copyright notice; 200 | keep intact all notices stating that this License and any 201 | non-permissive terms added in accord with section 7 apply to the code; 202 | keep intact all notices of the absence of any warranty; and give all 203 | recipients a copy of this License along with the Program. 204 | 205 | You may charge any price or no price for each copy that you convey, 206 | and you may offer support or warranty protection for a fee. 207 | 208 | 5. Conveying Modified Source Versions. 209 | 210 | You may convey a work based on the Program, or the modifications to 211 | produce it from the Program, in the form of source code under the 212 | terms of section 4, provided that you also meet all of these conditions: 213 | 214 | a) The work must carry prominent notices stating that you modified 215 | it, and giving a relevant date. 216 | 217 | b) The work must carry prominent notices stating that it is 218 | released under this License and any conditions added under section 219 | 7. This requirement modifies the requirement in section 4 to 220 | "keep intact all notices". 221 | 222 | c) You must license the entire work, as a whole, under this 223 | License to anyone who comes into possession of a copy. This 224 | License will therefore apply, along with any applicable section 7 225 | additional terms, to the whole of the work, and all its parts, 226 | regardless of how they are packaged. This License gives no 227 | permission to license the work in any other way, but it does not 228 | invalidate such permission if you have separately received it. 229 | 230 | d) If the work has interactive user interfaces, each must display 231 | Appropriate Legal Notices; however, if the Program has interactive 232 | interfaces that do not display Appropriate Legal Notices, your 233 | work need not make them do so. 234 | 235 | A compilation of a covered work with other separate and independent 236 | works, which are not by their nature extensions of the covered work, 237 | and which are not combined with it such as to form a larger program, 238 | in or on a volume of a storage or distribution medium, is called an 239 | "aggregate" if the compilation and its resulting copyright are not 240 | used to limit the access or legal rights of the compilation's users 241 | beyond what the individual works permit. Inclusion of a covered work 242 | in an aggregate does not cause this License to apply to the other 243 | parts of the aggregate. 244 | 245 | 6. Conveying Non-Source Forms. 246 | 247 | You may convey a covered work in object code form under the terms 248 | of sections 4 and 5, provided that you also convey the 249 | machine-readable Corresponding Source under the terms of this License, 250 | in one of these ways: 251 | 252 | a) Convey the object code in, or embodied in, a physical product 253 | (including a physical distribution medium), accompanied by the 254 | Corresponding Source fixed on a durable physical medium 255 | customarily used for software interchange. 256 | 257 | b) Convey the object code in, or embodied in, a physical product 258 | (including a physical distribution medium), accompanied by a 259 | written offer, valid for at least three years and valid for as 260 | long as you offer spare parts or customer support for that product 261 | model, to give anyone who possesses the object code either (1) a 262 | copy of the Corresponding Source for all the software in the 263 | product that is covered by this License, on a durable physical 264 | medium customarily used for software interchange, for a price no 265 | more than your reasonable cost of physically performing this 266 | conveying of source, or (2) access to copy the 267 | Corresponding Source from a network server at no charge. 268 | 269 | c) Convey individual copies of the object code with a copy of the 270 | written offer to provide the Corresponding Source. This 271 | alternative is allowed only occasionally and noncommercially, and 272 | only if you received the object code with such an offer, in accord 273 | with subsection 6b. 274 | 275 | d) Convey the object code by offering access from a designated 276 | place (gratis or for a charge), and offer equivalent access to the 277 | Corresponding Source in the same way through the same place at no 278 | further charge. You need not require recipients to copy the 279 | Corresponding Source along with the object code. If the place to 280 | copy the object code is a network server, the Corresponding Source 281 | may be on a different server (operated by you or a third party) 282 | that supports equivalent copying facilities, provided you maintain 283 | clear directions next to the object code saying where to find the 284 | Corresponding Source. Regardless of what server hosts the 285 | Corresponding Source, you remain obligated to ensure that it is 286 | available for as long as needed to satisfy these requirements. 287 | 288 | e) Convey the object code using peer-to-peer transmission, provided 289 | you inform other peers where the object code and Corresponding 290 | Source of the work are being offered to the general public at no 291 | charge under subsection 6d. 292 | 293 | A separable portion of the object code, whose source code is excluded 294 | from the Corresponding Source as a System Library, need not be 295 | included in conveying the object code work. 296 | 297 | A "User Product" is either (1) a "consumer product", which means any 298 | tangible personal property which is normally used for personal, family, 299 | or household purposes, or (2) anything designed or sold for incorporation 300 | into a dwelling. In determining whether a product is a consumer product, 301 | doubtful cases shall be resolved in favor of coverage. For a particular 302 | product received by a particular user, "normally used" refers to a 303 | typical or common use of that class of product, regardless of the status 304 | of the particular user or of the way in which the particular user 305 | actually uses, or expects or is expected to use, the product. A product 306 | is a consumer product regardless of whether the product has substantial 307 | commercial, industrial or non-consumer uses, unless such uses represent 308 | the only significant mode of use of the product. 309 | 310 | "Installation Information" for a User Product means any methods, 311 | procedures, authorization keys, or other information required to install 312 | and execute modified versions of a covered work in that User Product from 313 | a modified version of its Corresponding Source. The information must 314 | suffice to ensure that the continued functioning of the modified object 315 | code is in no case prevented or interfered with solely because 316 | modification has been made. 317 | 318 | If you convey an object code work under this section in, or with, or 319 | specifically for use in, a User Product, and the conveying occurs as 320 | part of a transaction in which the right of possession and use of the 321 | User Product is transferred to the recipient in perpetuity or for a 322 | fixed term (regardless of how the transaction is characterized), the 323 | Corresponding Source conveyed under this section must be accompanied 324 | by the Installation Information. But this requirement does not apply 325 | if neither you nor any third party retains the ability to install 326 | modified object code on the User Product (for example, the work has 327 | been installed in ROM). 328 | 329 | The requirement to provide Installation Information does not include a 330 | requirement to continue to provide support service, warranty, or updates 331 | for a work that has been modified or installed by the recipient, or for 332 | the User Product in which it has been modified or installed. Access to a 333 | network may be denied when the modification itself materially and 334 | adversely affects the operation of the network or violates the rules and 335 | protocols for communication across the network. 336 | 337 | Corresponding Source conveyed, and Installation Information provided, 338 | in accord with this section must be in a format that is publicly 339 | documented (and with an implementation available to the public in 340 | source code form), and must require no special password or key for 341 | unpacking, reading or copying. 342 | 343 | 7. Additional Terms. 344 | 345 | "Additional permissions" are terms that supplement the terms of this 346 | License by making exceptions from one or more of its conditions. 347 | Additional permissions that are applicable to the entire Program shall 348 | be treated as though they were included in this License, to the extent 349 | that they are valid under applicable law. If additional permissions 350 | apply only to part of the Program, that part may be used separately 351 | under those permissions, but the entire Program remains governed by 352 | this License without regard to the additional permissions. 353 | 354 | When you convey a copy of a covered work, you may at your option 355 | remove any additional permissions from that copy, or from any part of 356 | it. (Additional permissions may be written to require their own 357 | removal in certain cases when you modify the work.) You may place 358 | additional permissions on material, added by you to a covered work, 359 | for which you have or can give appropriate copyright permission. 360 | 361 | Notwithstanding any other provision of this License, for material you 362 | add to a covered work, you may (if authorized by the copyright holders of 363 | that material) supplement the terms of this License with terms: 364 | 365 | a) Disclaiming warranty or limiting liability differently from the 366 | terms of sections 15 and 16 of this License; or 367 | 368 | b) Requiring preservation of specified reasonable legal notices or 369 | author attributions in that material or in the Appropriate Legal 370 | Notices displayed by works containing it; or 371 | 372 | c) Prohibiting misrepresentation of the origin of that material, or 373 | requiring that modified versions of such material be marked in 374 | reasonable ways as different from the original version; or 375 | 376 | d) Limiting the use for publicity purposes of names of licensors or 377 | authors of the material; or 378 | 379 | e) Declining to grant rights under trademark law for use of some 380 | trade names, trademarks, or service marks; or 381 | 382 | f) Requiring indemnification of licensors and authors of that 383 | material by anyone who conveys the material (or modified versions of 384 | it) with contractual assumptions of liability to the recipient, for 385 | any liability that these contractual assumptions directly impose on 386 | those licensors and authors. 387 | 388 | All other non-permissive additional terms are considered "further 389 | restrictions" within the meaning of section 10. If the Program as you 390 | received it, or any part of it, contains a notice stating that it is 391 | governed by this License along with a term that is a further 392 | restriction, you may remove that term. If a license document contains 393 | a further restriction but permits relicensing or conveying under this 394 | License, you may add to a covered work material governed by the terms 395 | of that license document, provided that the further restriction does 396 | not survive such relicensing or conveying. 397 | 398 | If you add terms to a covered work in accord with this section, you 399 | must place, in the relevant source files, a statement of the 400 | additional terms that apply to those files, or a notice indicating 401 | where to find the applicable terms. 402 | 403 | Additional terms, permissive or non-permissive, may be stated in the 404 | form of a separately written license, or stated as exceptions; 405 | the above requirements apply either way. 406 | 407 | 8. Termination. 408 | 409 | You may not propagate or modify a covered work except as expressly 410 | provided under this License. Any attempt otherwise to propagate or 411 | modify it is void, and will automatically terminate your rights under 412 | this License (including any patent licenses granted under the third 413 | paragraph of section 11). 414 | 415 | However, if you cease all violation of this License, then your 416 | license from a particular copyright holder is reinstated (a) 417 | provisionally, unless and until the copyright holder explicitly and 418 | finally terminates your license, and (b) permanently, if the copyright 419 | holder fails to notify you of the violation by some reasonable means 420 | prior to 60 days after the cessation. 421 | 422 | Moreover, your license from a particular copyright holder is 423 | reinstated permanently if the copyright holder notifies you of the 424 | violation by some reasonable means, this is the first time you have 425 | received notice of violation of this License (for any work) from that 426 | copyright holder, and you cure the violation prior to 30 days after 427 | your receipt of the notice. 428 | 429 | Termination of your rights under this section does not terminate the 430 | licenses of parties who have received copies or rights from you under 431 | this License. If your rights have been terminated and not permanently 432 | reinstated, you do not qualify to receive new licenses for the same 433 | material under section 10. 434 | 435 | 9. Acceptance Not Required for Having Copies. 436 | 437 | You are not required to accept this License in order to receive or 438 | run a copy of the Program. Ancillary propagation of a covered work 439 | occurring solely as a consequence of using peer-to-peer transmission 440 | to receive a copy likewise does not require acceptance. However, 441 | nothing other than this License grants you permission to propagate or 442 | modify any covered work. These actions infringe copyright if you do 443 | not accept this License. Therefore, by modifying or propagating a 444 | covered work, you indicate your acceptance of this License to do so. 445 | 446 | 10. Automatic Licensing of Downstream Recipients. 447 | 448 | Each time you convey a covered work, the recipient automatically 449 | receives a license from the original licensors, to run, modify and 450 | propagate that work, subject to this License. You are not responsible 451 | for enforcing compliance by third parties with this License. 452 | 453 | An "entity transaction" is a transaction transferring control of an 454 | organization, or substantially all assets of one, or subdividing an 455 | organization, or merging organizations. If propagation of a covered 456 | work results from an entity transaction, each party to that 457 | transaction who receives a copy of the work also receives whatever 458 | licenses to the work the party's predecessor in interest had or could 459 | give under the previous paragraph, plus a right to possession of the 460 | Corresponding Source of the work from the predecessor in interest, if 461 | the predecessor has it or can get it with reasonable efforts. 462 | 463 | You may not impose any further restrictions on the exercise of the 464 | rights granted or affirmed under this License. For example, you may 465 | not impose a license fee, royalty, or other charge for exercise of 466 | rights granted under this License, and you may not initiate litigation 467 | (including a cross-claim or counterclaim in a lawsuit) alleging that 468 | any patent claim is infringed by making, using, selling, offering for 469 | sale, or importing the Program or any portion of it. 470 | 471 | 11. Patents. 472 | 473 | A "contributor" is a copyright holder who authorizes use under this 474 | License of the Program or a work on which the Program is based. The 475 | work thus licensed is called the contributor's "contributor version". 476 | 477 | A contributor's "essential patent claims" are all patent claims 478 | owned or controlled by the contributor, whether already acquired or 479 | hereafter acquired, that would be infringed by some manner, permitted 480 | by this License, of making, using, or selling its contributor version, 481 | but do not include claims that would be infringed only as a 482 | consequence of further modification of the contributor version. For 483 | purposes of this definition, "control" includes the right to grant 484 | patent sublicenses in a manner consistent with the requirements of 485 | this License. 486 | 487 | Each contributor grants you a non-exclusive, worldwide, royalty-free 488 | patent license under the contributor's essential patent claims, to 489 | make, use, sell, offer for sale, import and otherwise run, modify and 490 | propagate the contents of its contributor version. 491 | 492 | In the following three paragraphs, a "patent license" is any express 493 | agreement or commitment, however denominated, not to enforce a patent 494 | (such as an express permission to practice a patent or covenant not to 495 | sue for patent infringement). To "grant" such a patent license to a 496 | party means to make such an agreement or commitment not to enforce a 497 | patent against the party. 498 | 499 | If you convey a covered work, knowingly relying on a patent license, 500 | and the Corresponding Source of the work is not available for anyone 501 | to copy, free of charge and under the terms of this License, through a 502 | publicly available network server or other readily accessible means, 503 | then you must either (1) cause the Corresponding Source to be so 504 | available, or (2) arrange to deprive yourself of the benefit of the 505 | patent license for this particular work, or (3) arrange, in a manner 506 | consistent with the requirements of this License, to extend the patent 507 | license to downstream recipients. "Knowingly relying" means you have 508 | actual knowledge that, but for the patent license, your conveying the 509 | covered work in a country, or your recipient's use of the covered work 510 | in a country, would infringe one or more identifiable patents in that 511 | country that you have reason to believe are valid. 512 | 513 | If, pursuant to or in connection with a single transaction or 514 | arrangement, you convey, or propagate by procuring conveyance of, a 515 | covered work, and grant a patent license to some of the parties 516 | receiving the covered work authorizing them to use, propagate, modify 517 | or convey a specific copy of the covered work, then the patent license 518 | you grant is automatically extended to all recipients of the covered 519 | work and works based on it. 520 | 521 | A patent license is "discriminatory" if it does not include within 522 | the scope of its coverage, prohibits the exercise of, or is 523 | conditioned on the non-exercise of one or more of the rights that are 524 | specifically granted under this License. You may not convey a covered 525 | work if you are a party to an arrangement with a third party that is 526 | in the business of distributing software, under which you make payment 527 | to the third party based on the extent of your activity of conveying 528 | the work, and under which the third party grants, to any of the 529 | parties who would receive the covered work from you, a discriminatory 530 | patent license (a) in connection with copies of the covered work 531 | conveyed by you (or copies made from those copies), or (b) primarily 532 | for and in connection with specific products or compilations that 533 | contain the covered work, unless you entered into that arrangement, 534 | or that patent license was granted, prior to 28 March 2007. 535 | 536 | Nothing in this License shall be construed as excluding or limiting 537 | any implied license or other defenses to infringement that may 538 | otherwise be available to you under applicable patent law. 539 | 540 | 12. No Surrender of Others' Freedom. 541 | 542 | If conditions are imposed on you (whether by court order, agreement or 543 | otherwise) that contradict the conditions of this License, they do not 544 | excuse you from the conditions of this License. If you cannot convey a 545 | covered work so as to satisfy simultaneously your obligations under this 546 | License and any other pertinent obligations, then as a consequence you may 547 | not convey it at all. For example, if you agree to terms that obligate you 548 | to collect a royalty for further conveying from those to whom you convey 549 | the Program, the only way you could satisfy both those terms and this 550 | License would be to refrain entirely from conveying the Program. 551 | 552 | 13. Use with the GNU Affero General Public License. 553 | 554 | Notwithstanding any other provision of this License, you have 555 | permission to link or combine any covered work with a work licensed 556 | under version 3 of the GNU Affero General Public License into a single 557 | combined work, and to convey the resulting work. The terms of this 558 | License will continue to apply to the part which is the covered work, 559 | but the special requirements of the GNU Affero General Public License, 560 | section 13, concerning interaction through a network will apply to the 561 | combination as such. 562 | 563 | 14. Revised Versions of this License. 564 | 565 | The Free Software Foundation may publish revised and/or new versions of 566 | the GNU General Public License from time to time. Such new versions will 567 | be similar in spirit to the present version, but may differ in detail to 568 | address new problems or concerns. 569 | 570 | Each version is given a distinguishing version number. If the 571 | Program specifies that a certain numbered version of the GNU General 572 | Public License "or any later version" applies to it, you have the 573 | option of following the terms and conditions either of that numbered 574 | version or of any later version published by the Free Software 575 | Foundation. If the Program does not specify a version number of the 576 | GNU General Public License, you may choose any version ever published 577 | by the Free Software Foundation. 578 | 579 | If the Program specifies that a proxy can decide which future 580 | versions of the GNU General Public License can be used, that proxy's 581 | public statement of acceptance of a version permanently authorizes you 582 | to choose that version for the Program. 583 | 584 | Later license versions may give you additional or different 585 | permissions. However, no additional obligations are imposed on any 586 | author or copyright holder as a result of your choosing to follow a 587 | later version. 588 | 589 | 15. Disclaimer of Warranty. 590 | 591 | THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY 592 | APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT 593 | HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY 594 | OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, 595 | THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 596 | PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM 597 | IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF 598 | ALL NECESSARY SERVICING, REPAIR OR CORRECTION. 599 | 600 | 16. Limitation of Liability. 601 | 602 | IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING 603 | WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS 604 | THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY 605 | GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE 606 | USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF 607 | DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD 608 | PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), 609 | EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF 610 | SUCH DAMAGES. 611 | 612 | 17. Interpretation of Sections 15 and 16. 613 | 614 | If the disclaimer of warranty and limitation of liability provided 615 | above cannot be given local legal effect according to their terms, 616 | reviewing courts shall apply local law that most closely approximates 617 | an absolute waiver of all civil liability in connection with the 618 | Program, unless a warranty or assumption of liability accompanies a 619 | copy of the Program in return for a fee. 620 | 621 | END OF TERMS AND CONDITIONS 622 | 623 | How to Apply These Terms to Your New Programs 624 | 625 | If you develop a new program, and you want it to be of the greatest 626 | possible use to the public, the best way to achieve this is to make it 627 | free software which everyone can redistribute and change under these terms. 628 | 629 | To do so, attach the following notices to the program. It is safest 630 | to attach them to the start of each source file to most effectively 631 | state the exclusion of warranty; and each file should have at least 632 | the "copyright" line and a pointer to where the full notice is found. 633 | 634 | 635 | Copyright (C) 636 | 637 | This program is free software: you can redistribute it and/or modify 638 | it under the terms of the GNU General Public License as published by 639 | the Free Software Foundation, either version 3 of the License, or 640 | (at your option) any later version. 641 | 642 | This program is distributed in the hope that it will be useful, 643 | but WITHOUT ANY WARRANTY; without even the implied warranty of 644 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 645 | GNU General Public License for more details. 646 | 647 | You should have received a copy of the GNU General Public License 648 | along with this program. If not, see . 649 | 650 | Also add information on how to contact you by electronic and paper mail. 651 | 652 | If the program does terminal interaction, make it output a short 653 | notice like this when it starts in an interactive mode: 654 | 655 | Copyright (C) 656 | This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. 657 | This is free software, and you are welcome to redistribute it 658 | under certain conditions; type `show c' for details. 659 | 660 | The hypothetical commands `show w' and `show c' should show the appropriate 661 | parts of the General Public License. Of course, your program's commands 662 | might be different; for a GUI interface, you would use an "about box". 663 | 664 | You should also get your employer (if you work as a programmer) or school, 665 | if any, to sign a "copyright disclaimer" for the program, if necessary. 666 | For more information on this, and how to apply and follow the GNU GPL, see 667 | . 668 | 669 | The GNU General Public License does not permit incorporating your program 670 | into proprietary programs. If your program is a subroutine library, you 671 | may consider it more useful to permit linking proprietary applications with 672 | the library. If this is what you want to do, use the GNU Lesser General 673 | Public License instead of this License. But first, please read 674 | . -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 |   6 | 7 | ### This repo only focuses on NMS speed improvement based on https://github.com/ultralytics/yolov5. 8 | 9 | ### See `non_max_suppression` function of [utils/general.py](utils/general.py) for our Cluster-NMS implementation. 10 | 11 | # Batch mode Cluster-NMS 12 | 13 | Torchvision NMS has the fastest speed but fails to run in batch mode. 14 | 15 | Batch mode Cluster-NMS is made for this. 16 | 17 | ### Our goal is that when using TTA for getting better performance, NMS no longer becomes a potential time-consuming growth factor. 18 | 19 | ### Some Pretrained Weights 20 | 21 | | Model | APval | APtest | AP50 | SpeedGPU | FPSGPU || params | FLOPS | 22 | |---------- |------ |------ |------ | -------- | ------| ------ |------ | :------: | 23 | | [YOLOv5s](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 37.0 | 37.0 | 56.2 | **2.4ms** | **416** || 7.5M | 13.2B 24 | | [YOLOv5m](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 44.3 | 44.3 | 63.2 | 3.4ms | 294 || 21.8M | 39.4B 25 | | [YOLOv5l](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 47.7 | 47.7 | 66.5 | 4.4ms | 227 || 47.8M | 88.1B 26 | | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | **49.2** | **49.2** | **67.7** | 6.9ms | 145 || 89.0M | 166.4B 27 | | | | | | | || | 28 | | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/tag/v3.0) + TTA|**50.8**| **50.8** | **68.9** | 25.5ms | 39 || 89.0M | 354.3B 29 | | | | | | | || | 30 | | [YOLOv3-SPP](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 45.6 | 45.5 | 65.2 | 4.5ms | 222 || 63.0M | 118.0B 31 | 32 | For more details, please refer to https://github.com/ultralytics/yolov5. 33 | 34 | ## Requirements 35 | 36 | Python 3.8 or later with all [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) dependencies installed, including `torch>=1.6`. To install run: 37 | ```bash 38 | $ pip install -r requirements.txt 39 | ``` 40 | 41 | ## Evaluation for Batch Mode Weighted Cluster-NMS 42 | 43 | #### Hardware 44 | - 1 RTX 2080 Ti 45 | 46 | Evaluation command: `python test.py --weights yolov5s.pt --data coco.yaml --img 640 --augment --merge --batch-size 32` 47 | 48 | YOLOv5s.pt 49 | 50 | | NMS | TTA | max-box | weighted threshold | time (ms) | AP | AP50 | AP75 | APs | APm | APl | 51 | |:------------------------------------:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:| 52 | | Torchvision NMS | on | - | - | 3.2 / 17.9 | 38.0 | 56.5 | 41.2 | 20.9 | 42.6 | 51.7 | 53 | | Merge + Torchvision NMS | on | - | 0.65 | 3.2 / 18.6 | 38.0 | 56.5 | 41.4 | 20.9 | 42.7 | 51.8 | 54 | | Merge + Torchvision NMS | on | - | 0.8 | 3.2 / 18.9 | 38.1 | 56.5 | 41.4 | 21.0 | 42.7 | 51.8 | 55 | | Weighted Cluster-NMS | on | 1000 | 0.8 | 3.2 / 6.6 | 38.0 | 55.7 | 41.6 | 20.5 | 42.8 | 51.9 | 56 | | Weighted Cluster-NMS | on | 1500 | 0.65 | 3.2 / 10.2 | 38.1 | 56.1 | 41.9 | 20.9 | 42.7 | 51.8 | 57 | | Weighted Cluster-NMS | on | 1500 | 0.8 | 3.2 / 10.2 | 38.3 | 56.2 | 41.8 | 21.1 | 43.0 | 52.0 | 58 | | Weighted Cluster-NMS | on | 2000 | 0.8 | 3.2 / 14.5 | 38.4 | 56.4 | 41.9 | 21.3 | 43.1 | 52.1 | 59 | | | | | | | | | | | | | 60 | | Torchvision NMS | off | - | - | 1.5 / 5.4 | 36.9 | 56.2 | 40.0 | 21.0 | 42.1 | 47.4 | 61 | | Merge + Torchvision NMS | off | - | 0.65 | 1.3 / 6.7 | 36.9 | 56.2 | 40.2 | 20.9 | 42.1 | 47.4 | 62 | | Merge + Torchvision NMS | off | - | 0.8 | 1.3 / 6.7 | 37.1 | 56.2 | 40.3 | 21.1 | 42.2 | 47.6 | 63 | | Weighted Cluster-NMS | off | 1000 | 0.65 | 1.3 / 6.5 | 36.9 | 56.0 | 40.2 | 20.9 | 42.0 | 47.3 | 64 | | Weighted Cluster-NMS | off | 1000 | 0.8 | 1.3 / 6.5 | 37.0 | 56.0 | 40.3 | 21.1 | 42.2 | 47.5 | 65 | 66 | YOLOv5m.pt 67 | 68 | | NMS | TTA | max-box | weighted threshold | time (ms) | AP | AP50 | AP75 | APs | APm | APl | 69 | |:------------------------------------:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:| 70 | | Torchvision NMS | on | - | - | 6.4 / 10.4 | 45.1 | 63.2 | 49.0 | 27.0 | 50.2 | 60.5 | 71 | | Merge + Torchvision NMS | on | - | 0.65 | 6.4 / 11.5 | 45.0 | 63.2 | 49.0 | 26.9 | 50.2 | 60.3 | 72 | | Merge + Torchvision NMS | on | - | 0.8 | 6.4 / 11.5 | 45.2 | 63.3 | 49.1 | 27.0 | 50.3 | 60.5 | 73 | | Weighted Cluster-NMS | on | 1000 | 0.65 | 6.4 / 6.8 | 44.6 | 62.3 | 49.1 | 26.0 | 50.0 | 60.4 | 74 | | Weighted Cluster-NMS | on | 1500 | 0.65 | 6.4 / 9.8 | 44.9 | 62.9 | 49.4 | 26.6 | 50.2 | 60.4 | 75 | | Weighted Cluster-NMS | on | 1500 | 0.8 | 6.4 / 9.8 | 45.2 | 62.9 | 49.4 | 26.8 | 50.4 | 60.5 | 76 | | | | | | | | | | | | | 77 | | Torchvision NMS | off | - | - | 2.7 / 4.5 | 44.3 | 63.2 | 48.2 | 27.4 | 50.0 | 56.4 | 78 | | Merge + Torchvision NMS | off | - | 0.65 | 2.7 / 6.1 | 44.2 | 63.1 | 48.4 | 27.4 | 50.1 | 56.2 | 79 | | Merge + Torchvision NMS | off | - | 0.8 | 2.7 / 6.1 | 44.4 | 63.2 | 48.6 | 27.6 | 50.2 | 56.4 | 80 | | Weighted Cluster-NMS | off | 1000 | 0.65 | 2.7 / 6.1 | 44.2 | 62.9 | 48.5 | 27.3 | 50.0 | 56.3 | 81 | | Weighted Cluster-NMS | off | 1000 | 0.8 | 2.7 / 6.1 | 44.3 | 62.9 | 48.5 | 27.4 | 50.1 | 56.4 | 82 | 83 | YOLOv5x.pt `python test.py --weights yolov5s.pt --data coco.yaml --img 832 --augment --merge --batch-size 32` 84 | 85 | | NMS | TTA | max-box | weighted threshold | time (ms) | AP | AP50 | AP75 | APs | APm | APl | 86 | |:------------------------------------:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|:----:| 87 | | Merge + Torchvision NMS | on | - | 0.65 | 31.7 / 10.7 | 50.2 | 68.5 | 55.2 | 34.2 | 54.9 | 64.0 | 88 | | Weighted Cluster-NMS | on | 1500 | 0.8 | 31.8 / 9.9 | 50.3 | 68.0 | 55.4 | 33.9 | 55.1 | 64.6 | 89 | 90 | #### Details: 91 | - AP reports on `coco 2017val`. 92 | - `TTA` denotes Test-Time Augmentation. 93 | - `max-box` denotes maximum number of boxes processed in Batch Mode Cluster-NMS. 94 | - `weighted threshold` denotes the threshold used in weighted coordinates. 95 | - time reports model inference / NMS. 96 | - To avoid randomness, NMS runs three times here. See [test.py](test.py). 97 | ``` 98 | # Run NMS 99 | t = time_synchronized() 100 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge) 101 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge) 102 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge) 103 | t1 += time_synchronized() - t 104 | ``` 105 | 106 | ## Conclusion 107 | 108 | - Batch mode Weighted Cluster-NMS will have comparable speed with Torchvision merge NMS when batchsize>=16 and without TTA. 109 | - When using TTA, the time of torchvision NMS will increase significantly, because the model predicts much more boxes. Especially when using multi-scale testing or more TTA means. 110 | - Observed from experience, when using TTA, max-box = 1500 will be good. And when TTA is turned off, max-box = 1000. 111 | 112 | ## Related issues 113 | 114 | * [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303) 115 | * [Model Ensembling](https://github.com/ultralytics/yolov5/issues/318) 116 | * [INCREASING NMS SPEED](https://github.com/ultralytics/yolov3/issues/679) 117 | * [TESTING/INFERENCE AUGMENTATION](https://github.com/ultralytics/yolov3/issues/931) 118 | 119 | ## Environments 120 | 121 | YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled): 122 | 123 | - **Google Colab Notebook** with free GPU: Open In Colab 124 | - **Kaggle Notebook** with free GPU: [https://www.kaggle.com/ultralytics/yolov5](https://www.kaggle.com/ultralytics/yolov5) 125 | - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) 126 | - **Docker Image** https://hub.docker.com/r/ultralytics/yolov5. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) ![Docker Pulls](https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker) 127 | 128 | 129 | 130 | 131 | ## Citation 132 | 133 | [![DOI](https://zenodo.org/badge/264818686.svg)](https://zenodo.org/badge/latestdoi/264818686) 134 | 135 | 136 | #### This is the code for our paper: 137 | - [Distance-IoU Loss: Faster and Better Learning for Bounding Box Regression](https://arxiv.org/abs/1911.08287) 138 | - [Enhancing Geometric Factors into Model Learning and Inference for Object Detection and Instance Segmentation](http://arxiv.org/abs/2005.03572) 139 | 140 | ``` 141 | @Inproceedings{zheng2020diou, 142 | author = {Zheng, Zhaohui and Wang, Ping and Liu, Wei and Li, Jinze and Ye, Rongguang and Ren, Dongwei}, 143 | title = {Distance-IoU Loss: Faster and Better Learning for Bounding Box Regression}, 144 | booktitle = {The AAAI Conference on Artificial Intelligence (AAAI)}, 145 | year = {2020}, 146 | } 147 | 148 | @Article{zheng2021ciou, 149 | author = {Zheng, Zhaohui and Wang, Ping and Ren, Dongwei and Liu, Wei and Ye, Rongguang and Hu, Qinghua and Zuo, Wangmeng}, 150 | title = {Enhancing Geometric Factors in Model Learning and Inference for Object Detection and Instance Segmentation}, 151 | booktitle = {IEEE Transactions on Cybernetics}, 152 | year = {2021}, 153 | } 154 | ``` 155 | -------------------------------------------------------------------------------- /data/coco.yaml: -------------------------------------------------------------------------------- 1 | # COCO 2017 dataset http://cocodataset.org 2 | # Train command: python train.py --data coco.yaml 3 | # Default dataset location is next to /yolov5: 4 | # /parent_folder 5 | # /coco 6 | # /yolov5 7 | 8 | 9 | # download command/URL (optional) 10 | download: bash data/scripts/get_coco.sh 11 | 12 | # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] 13 | train: ../coco/train2017.txt # 118287 images 14 | val: ../coco/val2017.txt # 5000 images 15 | test: ../coco/test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794 16 | 17 | # number of classes 18 | nc: 80 19 | 20 | # class names 21 | names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 22 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 23 | 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 24 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 25 | 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 26 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 27 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 28 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 29 | 'hair drier', 'toothbrush'] 30 | 31 | # Print classes 32 | # with open('data/coco.yaml') as f: 33 | # d = yaml.load(f, Loader=yaml.FullLoader) # dict 34 | # for i, x in enumerate(d['names']): 35 | # print(i, x) 36 | -------------------------------------------------------------------------------- /data/coco128.yaml: -------------------------------------------------------------------------------- 1 | # COCO 2017 dataset http://cocodataset.org - first 128 training images 2 | # Train command: python train.py --data coco128.yaml 3 | # Default dataset location is next to /yolov5: 4 | # /parent_folder 5 | # /coco128 6 | # /yolov5 7 | 8 | 9 | # download command/URL (optional) 10 | download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 11 | 12 | # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] 13 | train: ../coco128/images/train2017/ # 128 images 14 | val: ../coco128/images/train2017/ # 128 images 15 | 16 | # number of classes 17 | nc: 80 18 | 19 | # class names 20 | names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 21 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 22 | 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 23 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 24 | 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 25 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 26 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 27 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 28 | 'hair drier', 'toothbrush'] 29 | -------------------------------------------------------------------------------- /data/hyp.finetune.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for VOC finetuning 2 | # python train.py --batch 64 --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50 3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 4 | 5 | 6 | # Hyperparameter Evolution Results 7 | # Generations: 306 8 | # P R mAP.5 mAP.5:.95 box obj cls 9 | # Metrics: 0.6 0.936 0.896 0.684 0.0115 0.00805 0.00146 10 | 11 | lr0: 0.0032 12 | lrf: 0.12 13 | momentum: 0.843 14 | weight_decay: 0.00036 15 | giou: 0.0296 16 | cls: 0.243 17 | cls_pw: 0.631 18 | obj: 0.301 19 | obj_pw: 0.911 20 | iou_t: 0.2 21 | anchor_t: 2.91 22 | # anchors: 3.63 23 | fl_gamma: 0.0 24 | hsv_h: 0.0138 25 | hsv_s: 0.664 26 | hsv_v: 0.464 27 | degrees: 0.373 28 | translate: 0.245 29 | scale: 0.898 30 | shear: 0.602 31 | perspective: 0.0 32 | flipud: 0.00856 33 | fliplr: 0.5 34 | mixup: 0.243 35 | -------------------------------------------------------------------------------- /data/hyp.scratch.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for COCO training from scratch 2 | # python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300 3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 4 | 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf) 8 | momentum: 0.937 # SGD momentum/Adam beta1 9 | weight_decay: 0.0005 # optimizer weight decay 5e-4 10 | giou: 0.05 # box loss gain 11 | cls: 0.5 # cls loss gain 12 | cls_pw: 1.0 # cls BCELoss positive_weight 13 | obj: 1.0 # obj loss gain (scale with pixels) 14 | obj_pw: 1.0 # obj BCELoss positive_weight 15 | iou_t: 0.20 # IoU training threshold 16 | anchor_t: 4.0 # anchor-multiple threshold 17 | # anchors: 0 # anchors per output grid (0 to ignore) 18 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 19 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 20 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 21 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 22 | degrees: 0.0 # image rotation (+/- deg) 23 | translate: 0.1 # image translation (+/- fraction) 24 | scale: 0.5 # image scale (+/- gain) 25 | shear: 0.0 # image shear (+/- deg) 26 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 27 | flipud: 0.0 # image flip up-down (probability) 28 | fliplr: 0.5 # image flip left-right (probability) 29 | mixup: 0.0 # image mixup (probability) 30 | -------------------------------------------------------------------------------- /data/scripts/get_coco.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # COCO 2017 dataset http://cocodataset.org 3 | # Download command: bash data/scripts/get_coco.sh 4 | # Train command: python train.py --data coco.yaml 5 | # Default dataset location is next to /yolov5: 6 | # /parent_folder 7 | # /coco 8 | # /yolov5 9 | 10 | # Download/unzip labels 11 | echo 'Downloading COCO 2017 labels ...' 12 | d='../' # unzip directory 13 | f='coco2017labels.zip' && curl -L https://github.com/ultralytics/yolov5/releases/download/v1.0/$f -o $f 14 | unzip -q $f -d $d && rm $f 15 | 16 | # Download/unzip images 17 | echo 'Downloading COCO 2017 images ...' 18 | d='../coco/images' # unzip directory 19 | f='train2017.zip' && curl http://images.cocodataset.org/zips/$f -o $f && unzip -q $f -d $d && rm $f # 19G, 118k images 20 | f='val2017.zip' && curl http://images.cocodataset.org/zips/$f -o $f && unzip -q $f -d $d && rm $f # 1G, 5k images 21 | # f='test2017.zip' && curl http://images.cocodataset.org/zips/$f -o $f && unzip -q $f -d $d && rm $f # 7G, 41k images 22 | -------------------------------------------------------------------------------- /data/scripts/get_voc.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/ 3 | # Download command: bash data/scripts/get_voc.sh 4 | # Train command: python train.py --data voc.yaml 5 | # Default dataset location is next to /yolov5: 6 | # /parent_folder 7 | # /VOC 8 | # /yolov5 9 | 10 | start=$(date +%s) 11 | 12 | # handle optional download dir 13 | if [ -z "$1" ]; then 14 | # navigate to ~/tmp 15 | echo "navigating to ../tmp/ ..." 16 | mkdir -p ../tmp 17 | cd ../tmp/ 18 | else 19 | # check if is valid directory 20 | if [ ! -d $1 ]; then 21 | echo $1 "is not a valid directory" 22 | exit 0 23 | fi 24 | echo "navigating to" $1 "..." 25 | cd $1 26 | fi 27 | 28 | echo "Downloading VOC2007 trainval ..." 29 | # Download data 30 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar 31 | echo "Downloading VOC2007 test data ..." 32 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar 33 | echo "Done downloading." 34 | 35 | # Extract data 36 | echo "Extracting trainval ..." 37 | tar -xf VOCtrainval_06-Nov-2007.tar 38 | echo "Extracting test ..." 39 | tar -xf VOCtest_06-Nov-2007.tar 40 | echo "removing tars ..." 41 | rm VOCtrainval_06-Nov-2007.tar 42 | rm VOCtest_06-Nov-2007.tar 43 | 44 | end=$(date +%s) 45 | runtime=$((end - start)) 46 | 47 | echo "Completed in" $runtime "seconds" 48 | 49 | start=$(date +%s) 50 | 51 | # handle optional download dir 52 | if [ -z "$1" ]; then 53 | # navigate to ~/tmp 54 | echo "navigating to ../tmp/ ..." 55 | mkdir -p ../tmp 56 | cd ../tmp/ 57 | else 58 | # check if is valid directory 59 | if [ ! -d $1 ]; then 60 | echo $1 "is not a valid directory" 61 | exit 0 62 | fi 63 | echo "navigating to" $1 "..." 64 | cd $1 65 | fi 66 | 67 | echo "Downloading VOC2012 trainval ..." 68 | # Download data 69 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar 70 | echo "Done downloading." 71 | 72 | # Extract data 73 | echo "Extracting trainval ..." 74 | tar -xf VOCtrainval_11-May-2012.tar 75 | echo "removing tar ..." 76 | rm VOCtrainval_11-May-2012.tar 77 | 78 | end=$(date +%s) 79 | runtime=$((end - start)) 80 | 81 | echo "Completed in" $runtime "seconds" 82 | 83 | cd ../tmp 84 | echo "Spliting dataset..." 85 | python3 - "$@" <train.txt 145 | cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt 146 | 147 | python3 - "$@" <= 1 87 | p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() 88 | else: 89 | p, s, im0 = path, '', im0s 90 | 91 | save_path = str(Path(out) / Path(p).name) 92 | txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '') 93 | s += '%gx%g ' % img.shape[2:] # print string 94 | gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh 95 | if det is not None and len(det): 96 | # Rescale boxes from img_size to im0 size 97 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() 98 | 99 | # Print results 100 | for c in det[:, -1].unique(): 101 | n = (det[:, -1] == c).sum() # detections per class 102 | s += '%g %ss, ' % (n, names[int(c)]) # add to string 103 | 104 | # Write results 105 | for *xyxy, conf, cls in reversed(det): 106 | if save_txt: # Write to file 107 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 108 | with open(txt_path + '.txt', 'a') as f: 109 | f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format 110 | 111 | if save_img or view_img: # Add bbox to image 112 | label = '%s %.2f' % (names[int(cls)], conf) 113 | plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) 114 | 115 | # Print time (inference + NMS) 116 | print('%sDone. (%.3fs)' % (s, t2 - t1)) 117 | 118 | # Stream results 119 | if view_img: 120 | cv2.imshow(p, im0) 121 | if cv2.waitKey(1) == ord('q'): # q to quit 122 | raise StopIteration 123 | 124 | # Save results (image with detections) 125 | if save_img: 126 | if dataset.mode == 'images': 127 | cv2.imwrite(save_path, im0) 128 | else: 129 | if vid_path != save_path: # new video 130 | vid_path = save_path 131 | if isinstance(vid_writer, cv2.VideoWriter): 132 | vid_writer.release() # release previous video writer 133 | 134 | fourcc = 'mp4v' # output video codec 135 | fps = vid_cap.get(cv2.CAP_PROP_FPS) 136 | w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 137 | h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 138 | vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) 139 | vid_writer.write(im0) 140 | 141 | if save_txt or save_img: 142 | print('Results saved to %s' % Path(out)) 143 | if platform.system() == 'Darwin' and not opt.update: # MacOS 144 | os.system('open ' + save_path) 145 | 146 | print('Done. (%.3fs)' % (time.time() - t0)) 147 | 148 | 149 | if __name__ == '__main__': 150 | parser = argparse.ArgumentParser() 151 | parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 152 | parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam 153 | parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder 154 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 155 | parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') 156 | parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') 157 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 158 | parser.add_argument('--view-img', action='store_true', help='display results') 159 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 160 | parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') 161 | parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') 162 | parser.add_argument('--augment', action='store_true', help='augmented inference') 163 | parser.add_argument('--update', action='store_true', help='update all models') 164 | opt = parser.parse_args() 165 | print(opt) 166 | 167 | with torch.no_grad(): 168 | if opt.update: # update all models (to fix SourceChangeWarning) 169 | for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 170 | detect() 171 | strip_optimizer(opt.weights) 172 | else: 173 | detect() 174 | -------------------------------------------------------------------------------- /hubconf.py: -------------------------------------------------------------------------------- 1 | """File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/ 2 | 3 | Usage: 4 | import torch 5 | model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80) 6 | """ 7 | 8 | dependencies = ['torch', 'yaml'] 9 | import os 10 | 11 | import torch 12 | 13 | from models.yolo import Model 14 | from utils.google_utils import attempt_download 15 | 16 | 17 | def create(name, pretrained, channels, classes): 18 | """Creates a specified YOLOv5 model 19 | 20 | Arguments: 21 | name (str): name of model, i.e. 'yolov5s' 22 | pretrained (bool): load pretrained weights into the model 23 | channels (int): number of input channels 24 | classes (int): number of model classes 25 | 26 | Returns: 27 | pytorch model 28 | """ 29 | config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path 30 | try: 31 | model = Model(config, channels, classes) 32 | if pretrained: 33 | ckpt = '%s.pt' % name # checkpoint filename 34 | attempt_download(ckpt) # download if not found locally 35 | state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32 36 | state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter 37 | model.load_state_dict(state_dict, strict=False) # load 38 | return model 39 | 40 | except Exception as e: 41 | help_url = 'https://github.com/ultralytics/yolov5/issues/36' 42 | s = 'Cache maybe be out of date, deleting cache and retrying may solve this. See %s for help.' % help_url 43 | raise Exception(s) from e 44 | 45 | 46 | def yolov5s(pretrained=False, channels=3, classes=80): 47 | """YOLOv5-small model from https://github.com/ultralytics/yolov5 48 | 49 | Arguments: 50 | pretrained (bool): load pretrained weights into the model, default=False 51 | channels (int): number of input channels, default=3 52 | classes (int): number of model classes, default=80 53 | 54 | Returns: 55 | pytorch model 56 | """ 57 | return create('yolov5s', pretrained, channels, classes) 58 | 59 | 60 | def yolov5m(pretrained=False, channels=3, classes=80): 61 | """YOLOv5-medium model from https://github.com/ultralytics/yolov5 62 | 63 | Arguments: 64 | pretrained (bool): load pretrained weights into the model, default=False 65 | channels (int): number of input channels, default=3 66 | classes (int): number of model classes, default=80 67 | 68 | Returns: 69 | pytorch model 70 | """ 71 | return create('yolov5m', pretrained, channels, classes) 72 | 73 | 74 | def yolov5l(pretrained=False, channels=3, classes=80): 75 | """YOLOv5-large model from https://github.com/ultralytics/yolov5 76 | 77 | Arguments: 78 | pretrained (bool): load pretrained weights into the model, default=False 79 | channels (int): number of input channels, default=3 80 | classes (int): number of model classes, default=80 81 | 82 | Returns: 83 | pytorch model 84 | """ 85 | return create('yolov5l', pretrained, channels, classes) 86 | 87 | 88 | def yolov5x(pretrained=False, channels=3, classes=80): 89 | """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5 90 | 91 | Arguments: 92 | pretrained (bool): load pretrained weights into the model, default=False 93 | channels (int): number of input channels, default=3 94 | classes (int): number of model classes, default=80 95 | 96 | Returns: 97 | pytorch model 98 | """ 99 | return create('yolov5x', pretrained, channels, classes) 100 | -------------------------------------------------------------------------------- /inference/images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zzh-tju/yolov5/fca5b7b7a85bcadac394f99c1805e54d9c20a21f/inference/images/bus.jpg -------------------------------------------------------------------------------- /inference/images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zzh-tju/yolov5/fca5b7b7a85bcadac394f99c1805e54d9c20a21f/inference/images/zidane.jpg -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zzh-tju/yolov5/fca5b7b7a85bcadac394f99c1805e54d9c20a21f/models/__init__.py -------------------------------------------------------------------------------- /models/common.py: -------------------------------------------------------------------------------- 1 | # This file contains modules common to various models 2 | import math 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | def autopad(k, p=None): # kernel, padding 9 | # Pad to 'same' 10 | if p is None: 11 | p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad 12 | return p 13 | 14 | 15 | def DWConv(c1, c2, k=1, s=1, act=True): 16 | # Depthwise convolution 17 | return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) 18 | 19 | 20 | class Conv(nn.Module): 21 | # Standard convolution 22 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 23 | super(Conv, self).__init__() 24 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) 25 | self.bn = nn.BatchNorm2d(c2) 26 | self.act = nn.Hardswish() if act else nn.Identity() 27 | 28 | def forward(self, x): 29 | return self.act(self.bn(self.conv(x))) 30 | 31 | def fuseforward(self, x): 32 | return self.act(self.conv(x)) 33 | 34 | 35 | class Bottleneck(nn.Module): 36 | # Standard bottleneck 37 | def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion 38 | super(Bottleneck, self).__init__() 39 | c_ = int(c2 * e) # hidden channels 40 | self.cv1 = Conv(c1, c_, 1, 1) 41 | self.cv2 = Conv(c_, c2, 3, 1, g=g) 42 | self.add = shortcut and c1 == c2 43 | 44 | def forward(self, x): 45 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 46 | 47 | 48 | class BottleneckCSP(nn.Module): 49 | # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks 50 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 51 | super(BottleneckCSP, self).__init__() 52 | c_ = int(c2 * e) # hidden channels 53 | self.cv1 = Conv(c1, c_, 1, 1) 54 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 55 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 56 | self.cv4 = Conv(2 * c_, c2, 1, 1) 57 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 58 | self.act = nn.LeakyReLU(0.1, inplace=True) 59 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) 60 | 61 | def forward(self, x): 62 | y1 = self.cv3(self.m(self.cv1(x))) 63 | y2 = self.cv2(x) 64 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 65 | 66 | 67 | class SPP(nn.Module): 68 | # Spatial pyramid pooling layer used in YOLOv3-SPP 69 | def __init__(self, c1, c2, k=(5, 9, 13)): 70 | super(SPP, self).__init__() 71 | c_ = c1 // 2 # hidden channels 72 | self.cv1 = Conv(c1, c_, 1, 1) 73 | self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) 74 | self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) 75 | 76 | def forward(self, x): 77 | x = self.cv1(x) 78 | return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) 79 | 80 | 81 | class Focus(nn.Module): 82 | # Focus wh information into c-space 83 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 84 | super(Focus, self).__init__() 85 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 86 | 87 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) 88 | return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) 89 | 90 | 91 | class Concat(nn.Module): 92 | # Concatenate a list of tensors along dimension 93 | def __init__(self, dimension=1): 94 | super(Concat, self).__init__() 95 | self.d = dimension 96 | 97 | def forward(self, x): 98 | return torch.cat(x, self.d) 99 | 100 | 101 | class Flatten(nn.Module): 102 | # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions 103 | @staticmethod 104 | def forward(x): 105 | return x.view(x.size(0), -1) 106 | 107 | 108 | class Classify(nn.Module): 109 | # Classification head, i.e. x(b,c1,20,20) to x(b,c2) 110 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups 111 | super(Classify, self).__init__() 112 | self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1) 113 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) # to x(b,c2,1,1) 114 | self.flat = Flatten() 115 | 116 | def forward(self, x): 117 | z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list 118 | return self.flat(self.conv(z)) # flatten to x(b,c2) 119 | -------------------------------------------------------------------------------- /models/experimental.py: -------------------------------------------------------------------------------- 1 | # This file contains experimental modules 2 | 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | 7 | from models.common import Conv, DWConv 8 | from utils.google_utils import attempt_download 9 | 10 | 11 | class CrossConv(nn.Module): 12 | # Cross Convolution Downsample 13 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 14 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 15 | super(CrossConv, self).__init__() 16 | c_ = int(c2 * e) # hidden channels 17 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 18 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 19 | self.add = shortcut and c1 == c2 20 | 21 | def forward(self, x): 22 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 23 | 24 | 25 | class C3(nn.Module): 26 | # Cross Convolution CSP 27 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 28 | super(C3, self).__init__() 29 | c_ = int(c2 * e) # hidden channels 30 | self.cv1 = Conv(c1, c_, 1, 1) 31 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 32 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 33 | self.cv4 = Conv(2 * c_, c2, 1, 1) 34 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 35 | self.act = nn.LeakyReLU(0.1, inplace=True) 36 | self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)]) 37 | 38 | def forward(self, x): 39 | y1 = self.cv3(self.m(self.cv1(x))) 40 | y2 = self.cv2(x) 41 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 42 | 43 | 44 | class Sum(nn.Module): 45 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 46 | def __init__(self, n, weight=False): # n: number of inputs 47 | super(Sum, self).__init__() 48 | self.weight = weight # apply weights boolean 49 | self.iter = range(n - 1) # iter object 50 | if weight: 51 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights 52 | 53 | def forward(self, x): 54 | y = x[0] # no weight 55 | if self.weight: 56 | w = torch.sigmoid(self.w) * 2 57 | for i in self.iter: 58 | y = y + x[i + 1] * w[i] 59 | else: 60 | for i in self.iter: 61 | y = y + x[i + 1] 62 | return y 63 | 64 | 65 | class GhostConv(nn.Module): 66 | # Ghost Convolution https://github.com/huawei-noah/ghostnet 67 | def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups 68 | super(GhostConv, self).__init__() 69 | c_ = c2 // 2 # hidden channels 70 | self.cv1 = Conv(c1, c_, k, s, g, act) 71 | self.cv2 = Conv(c_, c_, 5, 1, c_, act) 72 | 73 | def forward(self, x): 74 | y = self.cv1(x) 75 | return torch.cat([y, self.cv2(y)], 1) 76 | 77 | 78 | class GhostBottleneck(nn.Module): 79 | # Ghost Bottleneck https://github.com/huawei-noah/ghostnet 80 | def __init__(self, c1, c2, k, s): 81 | super(GhostBottleneck, self).__init__() 82 | c_ = c2 // 2 83 | self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw 84 | DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw 85 | GhostConv(c_, c2, 1, 1, act=False)) # pw-linear 86 | self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), 87 | Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() 88 | 89 | def forward(self, x): 90 | return self.conv(x) + self.shortcut(x) 91 | 92 | 93 | class MixConv2d(nn.Module): 94 | # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595 95 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): 96 | super(MixConv2d, self).__init__() 97 | groups = len(k) 98 | if equal_ch: # equal c_ per group 99 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices 100 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels 101 | else: # equal weight.numel() per group 102 | b = [c2] + [0] * groups 103 | a = np.eye(groups + 1, groups, k=-1) 104 | a -= np.roll(a, 1, axis=1) 105 | a *= np.array(k) ** 2 106 | a[0] = 1 107 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 108 | 109 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) 110 | self.bn = nn.BatchNorm2d(c2) 111 | self.act = nn.LeakyReLU(0.1, inplace=True) 112 | 113 | def forward(self, x): 114 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 115 | 116 | 117 | class Ensemble(nn.ModuleList): 118 | # Ensemble of models 119 | def __init__(self): 120 | super(Ensemble, self).__init__() 121 | 122 | def forward(self, x, augment=False): 123 | y = [] 124 | for module in self: 125 | y.append(module(x, augment)[0]) 126 | # y = torch.stack(y).max(0)[0] # max ensemble 127 | # y = torch.cat(y, 1) # nms ensemble 128 | y = torch.stack(y).mean(0) # mean ensemble 129 | return y, None # inference, train output 130 | 131 | 132 | def attempt_load(weights, map_location=None): 133 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 134 | model = Ensemble() 135 | for w in weights if isinstance(weights, list) else [weights]: 136 | attempt_download(w) 137 | model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model 138 | 139 | if len(model) == 1: 140 | return model[-1] # return model 141 | else: 142 | print('Ensemble created with %s\n' % weights) 143 | for k in ['names', 'stride']: 144 | setattr(model, k, getattr(model[-1], k)) 145 | return model # return ensemble 146 | -------------------------------------------------------------------------------- /models/export.py: -------------------------------------------------------------------------------- 1 | """Exports a YOLOv5 *.pt model to ONNX and TorchScript formats 2 | 3 | Usage: 4 | $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 5 | """ 6 | 7 | import argparse 8 | 9 | import torch 10 | import torch.nn as nn 11 | 12 | import models 13 | from models.experimental import attempt_load 14 | from utils.activations import Hardswish 15 | from utils.general import set_logging 16 | 17 | if __name__ == '__main__': 18 | parser = argparse.ArgumentParser() 19 | parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/ 20 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width 21 | parser.add_argument('--batch-size', type=int, default=1, help='batch size') 22 | opt = parser.parse_args() 23 | opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand 24 | print(opt) 25 | set_logging() 26 | 27 | # Input 28 | img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection 29 | 30 | # Load PyTorch model 31 | model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model 32 | 33 | # Update model 34 | for k, m in model.named_modules(): 35 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability 36 | if isinstance(m, models.common.Conv) and isinstance(m.act, nn.Hardswish): 37 | m.act = Hardswish() # assign activation 38 | # if isinstance(m, models.yolo.Detect): 39 | # m.forward = m.forward_export # assign forward (optional) 40 | model.model[-1].export = True # set Detect() layer export=True 41 | y = model(img) # dry run 42 | 43 | # TorchScript export 44 | try: 45 | print('\nStarting TorchScript export with torch %s...' % torch.__version__) 46 | f = opt.weights.replace('.pt', '.torchscript.pt') # filename 47 | ts = torch.jit.trace(model, img) 48 | ts.save(f) 49 | print('TorchScript export success, saved as %s' % f) 50 | except Exception as e: 51 | print('TorchScript export failure: %s' % e) 52 | 53 | # ONNX export 54 | try: 55 | import onnx 56 | 57 | print('\nStarting ONNX export with onnx %s...' % onnx.__version__) 58 | f = opt.weights.replace('.pt', '.onnx') # filename 59 | torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], 60 | output_names=['classes', 'boxes'] if y is None else ['output']) 61 | 62 | # Checks 63 | onnx_model = onnx.load(f) # load onnx model 64 | onnx.checker.check_model(onnx_model) # check onnx model 65 | # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model 66 | print('ONNX export success, saved as %s' % f) 67 | except Exception as e: 68 | print('ONNX export failure: %s' % e) 69 | 70 | # CoreML export 71 | try: 72 | import coremltools as ct 73 | 74 | print('\nStarting CoreML export with coremltools %s...' % ct.__version__) 75 | # convert model from torchscript and apply pixel scaling as per detect.py 76 | model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) 77 | f = opt.weights.replace('.pt', '.mlmodel') # filename 78 | model.save(f) 79 | print('CoreML export success, saved as %s' % f) 80 | except Exception as e: 81 | print('CoreML export failure: %s' % e) 82 | 83 | # Finish 84 | print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.') 85 | -------------------------------------------------------------------------------- /models/hub/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3-SPP head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, SPP, [512, [5, 9, 13]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /models/hub/yolov5-fpn.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, Bottleneck, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 6, BottleneckCSP, [1024]], # 9 25 | ] 26 | 27 | # YOLOv5 FPN head 28 | head: 29 | [[-1, 3, BottleneckCSP, [1024, False]], # 10 (P5/32-large) 30 | 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 3, BottleneckCSP, [512, False]], # 14 (P4/16-medium) 35 | 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 3, BottleneckCSP, [256, False]], # 18 (P3/8-small) 40 | 41 | [[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 42 | ] 43 | -------------------------------------------------------------------------------- /models/hub/yolov5-panet.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [116,90, 156,198, 373,326] # P5/32 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [10,13, 16,30, 33,23] # P3/8 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 PANet head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P5, P4, P3) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolo.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | import math 4 | from copy import deepcopy 5 | from pathlib import Path 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat 11 | from models.experimental import MixConv2d, CrossConv, C3 12 | from utils.general import check_anchor_order, make_divisible, check_file, set_logging 13 | from utils.torch_utils import ( 14 | time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, select_device) 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | class Detect(nn.Module): 20 | stride = None # strides computed during build 21 | export = False # onnx export 22 | 23 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 24 | super(Detect, self).__init__() 25 | self.nc = nc # number of classes 26 | self.no = nc + 5 # number of outputs per anchor 27 | self.nl = len(anchors) # number of detection layers 28 | self.na = len(anchors[0]) // 2 # number of anchors 29 | self.grid = [torch.zeros(1)] * self.nl # init grid 30 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 31 | self.register_buffer('anchors', a) # shape(nl,na,2) 32 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 33 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 34 | 35 | def forward(self, x): 36 | # x = x.copy() # for profiling 37 | z = [] # inference output 38 | self.training |= self.export 39 | for i in range(self.nl): 40 | x[i] = self.m[i](x[i]) # conv 41 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 42 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 43 | 44 | if not self.training: # inference 45 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 46 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 47 | 48 | y = x[i].sigmoid() 49 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 50 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 51 | z.append(y.view(bs, -1, self.no)) 52 | 53 | return x if self.training else (torch.cat(z, 1), x) 54 | 55 | @staticmethod 56 | def _make_grid(nx=20, ny=20): 57 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 58 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 59 | 60 | 61 | class Model(nn.Module): 62 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes 63 | super(Model, self).__init__() 64 | if isinstance(cfg, dict): 65 | self.yaml = cfg # model dict 66 | else: # is *.yaml 67 | import yaml # for torch hub 68 | self.yaml_file = Path(cfg).name 69 | with open(cfg) as f: 70 | self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict 71 | 72 | # Define model 73 | if nc and nc != self.yaml['nc']: 74 | print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)) 75 | self.yaml['nc'] = nc # override yaml value 76 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out 77 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 78 | 79 | # Build strides, anchors 80 | m = self.model[-1] # Detect() 81 | if isinstance(m, Detect): 82 | s = 128 # 2x min stride 83 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 84 | m.anchors /= m.stride.view(-1, 1, 1) 85 | check_anchor_order(m) 86 | self.stride = m.stride 87 | self._initialize_biases() # only run once 88 | # print('Strides: %s' % m.stride.tolist()) 89 | 90 | # Init weights, biases 91 | initialize_weights(self) 92 | self.info() 93 | print('') 94 | 95 | def forward(self, x, augment=False, profile=False): 96 | if augment: 97 | img_size = x.shape[-2:] # height, width 98 | s = [1, 0.83, 0.67] # scales 99 | f = [None, 3, None] # flips (2-ud, 3-lr) 100 | y = [] # outputs 101 | for si, fi in zip(s, f): 102 | xi = scale_img(x.flip(fi) if fi else x, si) 103 | yi = self.forward_once(xi)[0] # forward 104 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 105 | yi[..., :4] /= si # de-scale 106 | if fi == 2: 107 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 108 | elif fi == 3: 109 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 110 | y.append(yi) 111 | return torch.cat(y, 1), None # augmented inference, train 112 | else: 113 | return self.forward_once(x, profile) # single-scale inference, train 114 | 115 | def forward_once(self, x, profile=False): 116 | y, dt = [], [] # outputs 117 | for m in self.model: 118 | if m.f != -1: # if not from previous layer 119 | x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers 120 | 121 | if profile: 122 | try: 123 | import thop 124 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS 125 | except: 126 | o = 0 127 | t = time_synchronized() 128 | for _ in range(10): 129 | _ = m(x) 130 | dt.append((time_synchronized() - t) * 100) 131 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 132 | 133 | x = m(x) # run 134 | y.append(x if m.i in self.save else None) # save output 135 | 136 | if profile: 137 | print('%.1fms total' % sum(dt)) 138 | return x 139 | 140 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 141 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 142 | m = self.model[-1] # Detect() module 143 | for mi, s in zip(m.m, m.stride): # from 144 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 145 | b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 146 | b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 147 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 148 | 149 | def _print_biases(self): 150 | m = self.model[-1] # Detect() module 151 | for mi in m.m: # from 152 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 153 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 154 | 155 | # def _print_weights(self): 156 | # for m in self.model.modules(): 157 | # if type(m) is Bottleneck: 158 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 159 | 160 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 161 | print('Fusing layers... ') 162 | for m in self.model.modules(): 163 | if type(m) is Conv: 164 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability 165 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 166 | delattr(m, 'bn') # remove batchnorm 167 | m.forward = m.fuseforward # update forward 168 | self.info() 169 | return self 170 | 171 | def info(self, verbose=False): # print model information 172 | model_info(self, verbose) 173 | 174 | 175 | def parse_model(d, ch): # model_dict, input_channels(3) 176 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 177 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 178 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 179 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 180 | 181 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 182 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 183 | m = eval(m) if isinstance(m, str) else m # eval strings 184 | for j, a in enumerate(args): 185 | try: 186 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 187 | except: 188 | pass 189 | 190 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 191 | if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: 192 | c1, c2 = ch[f], args[0] 193 | 194 | # Normal 195 | # if i > 0 and args[0] != no: # channel expansion factor 196 | # ex = 1.75 # exponential (default 2.0) 197 | # e = math.log(c2 / ch[1]) / math.log(2) 198 | # c2 = int(ch[1] * ex ** e) 199 | # if m != Focus: 200 | 201 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 202 | 203 | # Experimental 204 | # if i > 0 and args[0] != no: # channel expansion factor 205 | # ex = 1 + gw # exponential (default 2.0) 206 | # ch1 = 32 # ch[1] 207 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n 208 | # c2 = int(ch1 * ex ** e) 209 | # if m != Focus: 210 | # c2 = make_divisible(c2, 8) if c2 != no else c2 211 | 212 | args = [c1, c2, *args[1:]] 213 | if m in [BottleneckCSP, C3]: 214 | args.insert(2, n) 215 | n = 1 216 | elif m is nn.BatchNorm2d: 217 | args = [ch[f]] 218 | elif m is Concat: 219 | c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) 220 | elif m is Detect: 221 | args.append([ch[x + 1] for x in f]) 222 | if isinstance(args[1], int): # number of anchors 223 | args[1] = [list(range(args[1] * 2))] * len(f) 224 | else: 225 | c2 = ch[f] 226 | 227 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 228 | t = str(m)[8:-2].replace('__main__.', '') # module type 229 | np = sum([x.numel() for x in m_.parameters()]) # number params 230 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 231 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print 232 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 233 | layers.append(m_) 234 | ch.append(c2) 235 | return nn.Sequential(*layers), sorted(save) 236 | 237 | 238 | if __name__ == '__main__': 239 | parser = argparse.ArgumentParser() 240 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 241 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 242 | opt = parser.parse_args() 243 | opt.cfg = check_file(opt.cfg) # check file 244 | set_logging() 245 | device = select_device(opt.device) 246 | 247 | # Create model 248 | model = Model(opt.cfg).to(device) 249 | model.train() 250 | 251 | # Profile 252 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 253 | # y = model(img, profile=True) 254 | 255 | # ONNX export 256 | # model.model[-1].export = True 257 | # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) 258 | 259 | # Tensorboard 260 | # from torch.utils.tensorboard import SummaryWriter 261 | # tb_writer = SummaryWriter() 262 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") 263 | # tb_writer.add_graph(model.model, img) # add model to tensorboard 264 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard 265 | -------------------------------------------------------------------------------- /models/yolov5l.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5m.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.67 # model depth multiple 4 | width_multiple: 0.75 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5s.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.33 # model depth multiple 4 | width_multiple: 0.50 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5x.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.33 # model depth multiple 4 | width_multiple: 1.25 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # base ---------------------------------------- 4 | Cython 5 | matplotlib>=3.2.2 6 | numpy>=1.18.5 7 | opencv-python>=4.1.2 8 | pillow 9 | PyYAML>=5.3 10 | scipy>=1.4.1 11 | tensorboard>=2.2 12 | torch>=1.6.0 13 | torchvision>=0.7.0 14 | tqdm>=4.41.0 15 | 16 | # coco ---------------------------------------- 17 | # pycocotools>=2.0 18 | 19 | # export -------------------------------------- 20 | # packaging # for coremltools 21 | # coremltools==4.0b3 22 | # onnx>=1.7.0 23 | # scikit-learn==0.19.2 # for coreml quantization 24 | 25 | # extras -------------------------------------- 26 | # thop # FLOPS computation 27 | # seaborn # plotting 28 | -------------------------------------------------------------------------------- /sotabench.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import glob 3 | import json 4 | import os 5 | import shutil 6 | from pathlib import Path 7 | 8 | import numpy as np 9 | import torch 10 | import yaml 11 | from tqdm import tqdm 12 | 13 | from models.experimental import attempt_load 14 | from utils.datasets import create_dataloader 15 | from utils.general import ( 16 | coco80_to_coco91_class, check_dataset, check_file, check_img_size, compute_loss, non_max_suppression, scale_coords, 17 | xyxy2xywh, clip_coords, plot_images, xywh2xyxy, box_iou, output_to_target, ap_per_class, set_logging) 18 | from utils.torch_utils import select_device, time_synchronized 19 | 20 | 21 | from sotabencheval.object_detection import COCOEvaluator 22 | from sotabencheval.utils import is_server 23 | 24 | DATA_ROOT = './.data/vision/coco' if is_server() else '../coco' # sotabench data dir 25 | 26 | 27 | def test(data, 28 | weights=None, 29 | batch_size=16, 30 | imgsz=640, 31 | conf_thres=0.001, 32 | iou_thres=0.6, # for NMS 33 | save_json=False, 34 | single_cls=False, 35 | augment=False, 36 | verbose=False, 37 | model=None, 38 | dataloader=None, 39 | save_dir='', 40 | merge=False, 41 | save_txt=False): 42 | # Initialize/load model and set device 43 | training = model is not None 44 | if training: # called by train.py 45 | device = next(model.parameters()).device # get model device 46 | 47 | else: # called directly 48 | set_logging() 49 | device = select_device(opt.device, batch_size=batch_size) 50 | merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels 51 | if save_txt: 52 | out = Path('inference/output') 53 | if os.path.exists(out): 54 | shutil.rmtree(out) # delete output folder 55 | os.makedirs(out) # make new output folder 56 | 57 | # Remove previous 58 | for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')): 59 | os.remove(f) 60 | 61 | # Load model 62 | model = attempt_load(weights, map_location=device) # load FP32 model 63 | imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size 64 | 65 | # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99 66 | # if device.type != 'cpu' and torch.cuda.device_count() > 1: 67 | # model = nn.DataParallel(model) 68 | 69 | # Half 70 | half = device.type != 'cpu' # half precision only supported on CUDA 71 | if half: 72 | model.half() 73 | 74 | # Configure 75 | model.eval() 76 | with open(data) as f: 77 | data = yaml.load(f, Loader=yaml.FullLoader) # model dict 78 | check_dataset(data) # check 79 | nc = 1 if single_cls else int(data['nc']) # number of classes 80 | iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 81 | niou = iouv.numel() 82 | 83 | # Dataloader 84 | if not training: 85 | img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img 86 | _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once 87 | path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images 88 | dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, 89 | hyp=None, augment=False, cache=True, pad=0.5, rect=True)[0] 90 | 91 | seen = 0 92 | names = model.names if hasattr(model, 'names') else model.module.names 93 | coco91class = coco80_to_coco91_class() 94 | s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95') 95 | p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0. 96 | loss = torch.zeros(3, device=device) 97 | jdict, stats, ap, ap_class = [], [], [], [] 98 | evaluator = COCOEvaluator(root=DATA_ROOT, model_name=opt.weights.replace('.pt', '')) 99 | for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)): 100 | img = img.to(device, non_blocking=True) 101 | img = img.half() if half else img.float() # uint8 to fp16/32 102 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 103 | targets = targets.to(device) 104 | nb, _, height, width = img.shape # batch size, channels, height, width 105 | whwh = torch.Tensor([width, height, width, height]).to(device) 106 | 107 | # Disable gradients 108 | with torch.no_grad(): 109 | # Run model 110 | t = time_synchronized() 111 | inf_out, train_out = model(img, augment=augment) # inference and training outputs 112 | t0 += time_synchronized() - t 113 | 114 | # Compute loss 115 | if training: # if model has loss hyperparameters 116 | loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls 117 | 118 | # Run NMS 119 | t = time_synchronized() 120 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, merge=merge) 121 | t1 += time_synchronized() - t 122 | 123 | # Statistics per image 124 | for si, pred in enumerate(output): 125 | labels = targets[targets[:, 0] == si, 1:] 126 | nl = len(labels) 127 | tcls = labels[:, 0].tolist() if nl else [] # target class 128 | seen += 1 129 | 130 | if pred is None: 131 | if nl: 132 | stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls)) 133 | continue 134 | 135 | # Append to text file 136 | if save_txt: 137 | gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh 138 | x = pred.clone() 139 | x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4], shapes[si][0], shapes[si][1]) # to original 140 | for *xyxy, conf, cls in x: 141 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 142 | with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f: 143 | f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format 144 | 145 | # Clip boxes to image bounds 146 | clip_coords(pred, (height, width)) 147 | 148 | # Append to pycocotools JSON dictionary 149 | if save_json: 150 | # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... 151 | image_id = Path(paths[si]).stem 152 | box = pred[:, :4].clone() # xyxy 153 | scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape 154 | box = xyxy2xywh(box) # xywh 155 | box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner 156 | for p, b in zip(pred.tolist(), box.tolist()): 157 | result = {'image_id': int(image_id) if image_id.isnumeric() else image_id, 158 | 'category_id': coco91class[int(p[5])], 159 | 'bbox': [round(x, 3) for x in b], 160 | 'score': round(p[4], 5)} 161 | jdict.append(result) 162 | 163 | #evaluator.add([result]) 164 | #if evaluator.cache_exists: 165 | # break 166 | 167 | # # Assign all predictions as incorrect 168 | # correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device) 169 | # if nl: 170 | # detected = [] # target indices 171 | # tcls_tensor = labels[:, 0] 172 | # 173 | # # target boxes 174 | # tbox = xywh2xyxy(labels[:, 1:5]) * whwh 175 | # 176 | # # Per target class 177 | # for cls in torch.unique(tcls_tensor): 178 | # ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices 179 | # pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices 180 | # 181 | # # Search for detections 182 | # if pi.shape[0]: 183 | # # Prediction to target ious 184 | # ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices 185 | # 186 | # # Append detections 187 | # detected_set = set() 188 | # for j in (ious > iouv[0]).nonzero(as_tuple=False): 189 | # d = ti[i[j]] # detected target 190 | # if d.item() not in detected_set: 191 | # detected_set.add(d.item()) 192 | # detected.append(d) 193 | # correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn 194 | # if len(detected) == nl: # all targets already located in image 195 | # break 196 | # 197 | # # Append statistics (correct, conf, pcls, tcls) 198 | # stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) 199 | 200 | # # Plot images 201 | # if batch_i < 1: 202 | # f = Path(save_dir) / ('test_batch%g_gt.jpg' % batch_i) # filename 203 | # plot_images(img, targets, paths, str(f), names) # ground truth 204 | # f = Path(save_dir) / ('test_batch%g_pred.jpg' % batch_i) 205 | # plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions 206 | 207 | evaluator.add(jdict) 208 | evaluator.save() 209 | 210 | # # Compute statistics 211 | # stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy 212 | # if len(stats) and stats[0].any(): 213 | # p, r, ap, f1, ap_class = ap_per_class(*stats) 214 | # p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95] 215 | # mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() 216 | # nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class 217 | # else: 218 | # nt = torch.zeros(1) 219 | # 220 | # # Print results 221 | # pf = '%20s' + '%12.3g' * 6 # print format 222 | # print(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) 223 | # 224 | # # Print results per class 225 | # if verbose and nc > 1 and len(stats): 226 | # for i, c in enumerate(ap_class): 227 | # print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) 228 | # 229 | # # Print speeds 230 | # t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple 231 | # if not training: 232 | # print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t) 233 | # 234 | # # Save JSON 235 | # if save_json and len(jdict): 236 | # f = 'detections_val2017_%s_results.json' % \ 237 | # (weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename 238 | # print('\nCOCO mAP with pycocotools... saving %s...' % f) 239 | # with open(f, 'w') as file: 240 | # json.dump(jdict, file) 241 | # 242 | # try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb 243 | # from pycocotools.coco import COCO 244 | # from pycocotools.cocoeval import COCOeval 245 | # 246 | # imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] 247 | # cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api 248 | # cocoDt = cocoGt.loadRes(f) # initialize COCO pred api 249 | # cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') 250 | # cocoEval.params.imgIds = imgIds # image IDs to evaluate 251 | # cocoEval.evaluate() 252 | # cocoEval.accumulate() 253 | # cocoEval.summarize() 254 | # map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) 255 | # except Exception as e: 256 | # print('ERROR: pycocotools unable to run: %s' % e) 257 | # 258 | # # Return results 259 | # model.float() # for training 260 | # maps = np.zeros(nc) + map 261 | # for i, c in enumerate(ap_class): 262 | # maps[c] = ap[i] 263 | # return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t 264 | 265 | 266 | if __name__ == '__main__': 267 | parser = argparse.ArgumentParser(prog='test.py') 268 | parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 269 | parser.add_argument('--data', type=str, default='data/coco.yaml', help='*.data path') 270 | parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') 271 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 272 | parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') 273 | parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS') 274 | parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') 275 | parser.add_argument('--task', default='val', help="'val', 'test', 'study'") 276 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 277 | parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') 278 | parser.add_argument('--augment', action='store_true', help='augmented inference') 279 | parser.add_argument('--merge', action='store_true', help='use Merge NMS') 280 | parser.add_argument('--verbose', action='store_true', help='report mAP by class') 281 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 282 | opt = parser.parse_args() 283 | opt.save_json |= opt.data.endswith('coco.yaml') 284 | opt.data = check_file(opt.data) # check file 285 | print(opt) 286 | 287 | if opt.task in ['val', 'test']: # run normally 288 | test(opt.data, 289 | opt.weights, 290 | opt.batch_size, 291 | opt.img_size, 292 | opt.conf_thres, 293 | opt.iou_thres, 294 | opt.save_json, 295 | opt.single_cls, 296 | opt.augment, 297 | opt.verbose) 298 | 299 | elif opt.task == 'study': # run over a range of settings and save/plot 300 | for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 301 | f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to 302 | x = list(range(320, 800, 64)) # x axis 303 | y = [] # y axis 304 | for i in x: # img-size 305 | print('\nRunning %s point %s...' % (f, i)) 306 | r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json) 307 | y.append(r + t) # results and times 308 | np.savetxt(f, y, fmt='%10.4g') # save 309 | os.system('zip -r study.zip study_*.txt') 310 | # utils.general.plot_study_txt(f, x) # plot -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import glob 3 | import json 4 | import os 5 | import shutil 6 | from pathlib import Path 7 | 8 | import numpy as np 9 | import torch 10 | import yaml 11 | from tqdm import tqdm 12 | 13 | from models.experimental import attempt_load 14 | from utils.datasets import create_dataloader 15 | from utils.general import ( 16 | coco80_to_coco91_class, check_dataset, check_file, check_img_size, compute_loss, non_max_suppression, scale_coords, 17 | xyxy2xywh, clip_coords, plot_images, xywh2xyxy, box_iou, output_to_target, ap_per_class, set_logging) 18 | from utils.torch_utils import select_device, time_synchronized 19 | 20 | 21 | def test(data, 22 | weights=None, 23 | batch_size=16, 24 | max_box=1500, 25 | imgsz=640, 26 | conf_thres=0.001, 27 | iou_thres=0.6, # for NMS 28 | save_json=False, 29 | single_cls=False, 30 | augment=False, 31 | verbose=False, 32 | model=None, 33 | dataloader=None, 34 | save_dir='', 35 | merge=False, 36 | save_txt=False): 37 | # Initialize/load model and set device 38 | training = model is not None 39 | if training: # called by train.py 40 | device = next(model.parameters()).device # get model device 41 | 42 | else: # called directly 43 | set_logging() 44 | device = select_device(opt.device, batch_size=batch_size) 45 | merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels 46 | if save_txt: 47 | out = Path('inference/output') 48 | if os.path.exists(out): 49 | shutil.rmtree(out) # delete output folder 50 | os.makedirs(out) # make new output folder 51 | 52 | # Remove previous 53 | for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')): 54 | os.remove(f) 55 | 56 | # Load model 57 | model = attempt_load(weights, map_location=device) # load FP32 model 58 | imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size 59 | 60 | # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99 61 | # if device.type != 'cpu' and torch.cuda.device_count() > 1: 62 | # model = nn.DataParallel(model) 63 | 64 | # Half 65 | half = device.type != 'cpu' # half precision only supported on CUDA 66 | if half: 67 | model.half() 68 | 69 | # Configure 70 | model.eval() 71 | with open(data) as f: 72 | data = yaml.load(f, Loader=yaml.FullLoader) # model dict 73 | check_dataset(data) # check 74 | nc = 1 if single_cls else int(data['nc']) # number of classes 75 | iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 76 | niou = iouv.numel() 77 | 78 | # Dataloader 79 | if not training: 80 | img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img 81 | _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once 82 | path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images 83 | dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, 84 | hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0] 85 | 86 | seen = 0 87 | names = model.names if hasattr(model, 'names') else model.module.names 88 | coco91class = coco80_to_coco91_class() 89 | s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95') 90 | p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0. 91 | loss = torch.zeros(3, device=device) 92 | jdict, stats, ap, ap_class = [], [], [], [] 93 | for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)): 94 | img = img.to(device, non_blocking=True) 95 | img = img.half() if half else img.float() # uint8 to fp16/32 96 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 97 | targets = targets.to(device) 98 | nb, _, height, width = img.shape # batch size, channels, height, width 99 | whwh = torch.Tensor([width, height, width, height]).to(device) 100 | 101 | # Disable gradients 102 | with torch.no_grad(): 103 | # Run model 104 | t = time_synchronized() 105 | inf_out, train_out = model(img, augment=augment) # inference and training outputs 106 | t0 += time_synchronized() - t 107 | 108 | # Compute loss 109 | if training: # if model has loss hyperparameters 110 | loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls 111 | 112 | # Run NMS 113 | t = time_synchronized() 114 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge) 115 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge) 116 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge) 117 | t1 += time_synchronized() - t 118 | 119 | # Statistics per image 120 | for si, pred in enumerate(output): 121 | labels = targets[targets[:, 0] == si, 1:] 122 | nl = len(labels) 123 | tcls = labels[:, 0].tolist() if nl else [] # target class 124 | seen += 1 125 | 126 | if pred is None: 127 | if nl: 128 | stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls)) 129 | continue 130 | 131 | # Append to text file 132 | if save_txt: 133 | gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh 134 | x = pred.clone() 135 | x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4], shapes[si][0], shapes[si][1]) # to original 136 | for *xyxy, conf, cls in x: 137 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 138 | with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f: 139 | f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format 140 | 141 | # Clip boxes to image bounds 142 | clip_coords(pred, (height, width)) 143 | 144 | # Append to pycocotools JSON dictionary 145 | if save_json: 146 | # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... 147 | image_id = Path(paths[si]).stem 148 | box = pred[:, :4].clone() # xyxy 149 | scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape 150 | box = xyxy2xywh(box) # xywh 151 | box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner 152 | for p, b in zip(pred.tolist(), box.tolist()): 153 | jdict.append({'image_id': int(image_id) if image_id.isnumeric() else image_id, 154 | 'category_id': coco91class[int(p[5])], 155 | 'bbox': [round(x, 3) for x in b], 156 | 'score': round(p[4], 5)}) 157 | 158 | # Assign all predictions as incorrect 159 | correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device) 160 | if nl: 161 | detected = [] # target indices 162 | tcls_tensor = labels[:, 0] 163 | 164 | # target boxes 165 | tbox = xywh2xyxy(labels[:, 1:5]) * whwh 166 | 167 | # Per target class 168 | for cls in torch.unique(tcls_tensor): 169 | ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices 170 | pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices 171 | 172 | # Search for detections 173 | if pi.shape[0]: 174 | # Prediction to target ious 175 | ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices 176 | 177 | # Append detections 178 | detected_set = set() 179 | for j in (ious > iouv[0]).nonzero(as_tuple=False): 180 | d = ti[i[j]] # detected target 181 | if d.item() not in detected_set: 182 | detected_set.add(d.item()) 183 | detected.append(d) 184 | correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn 185 | if len(detected) == nl: # all targets already located in image 186 | break 187 | 188 | # Append statistics (correct, conf, pcls, tcls) 189 | stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) 190 | 191 | # Plot images 192 | if batch_i < 1: 193 | f = Path(save_dir) / ('test_batch%g_gt.jpg' % batch_i) # filename 194 | plot_images(img, targets, paths, str(f), names) # ground truth 195 | f = Path(save_dir) / ('test_batch%g_pred.jpg' % batch_i) 196 | plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions 197 | 198 | # Compute statistics 199 | stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy 200 | if len(stats) and stats[0].any(): 201 | p, r, ap, f1, ap_class = ap_per_class(*stats) 202 | p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95] 203 | mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() 204 | nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class 205 | else: 206 | nt = torch.zeros(1) 207 | 208 | # Print results 209 | pf = '%20s' + '%12.3g' * 6 # print format 210 | print(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) 211 | 212 | # Print results per class 213 | if verbose and nc > 1 and len(stats): 214 | for i, c in enumerate(ap_class): 215 | print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) 216 | 217 | # Print speeds 218 | t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple 219 | if not training: 220 | print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t) 221 | 222 | # Save JSON 223 | if save_json and len(jdict): 224 | f = 'detections_val2017_%s_results.json' % \ 225 | (weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename 226 | print('\nCOCO mAP with pycocotools... saving %s...' % f) 227 | with open(f, 'w') as file: 228 | json.dump(jdict, file) 229 | 230 | try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb 231 | from pycocotools.coco import COCO 232 | from pycocotools.cocoeval import COCOeval 233 | 234 | imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] 235 | cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api 236 | cocoDt = cocoGt.loadRes(f) # initialize COCO pred api 237 | cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') 238 | cocoEval.params.imgIds = imgIds # image IDs to evaluate 239 | cocoEval.evaluate() 240 | cocoEval.accumulate() 241 | cocoEval.summarize() 242 | map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) 243 | except Exception as e: 244 | print('ERROR: pycocotools unable to run: %s' % e) 245 | 246 | # Return results 247 | model.float() # for training 248 | maps = np.zeros(nc) + map 249 | for i, c in enumerate(ap_class): 250 | maps[c] = ap[i] 251 | return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t 252 | 253 | 254 | if __name__ == '__main__': 255 | parser = argparse.ArgumentParser(prog='test.py') 256 | parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 257 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path') 258 | parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') 259 | parser.add_argument('--max-box', type=int, default=1500, help='max number of boxes for cluster_nms') 260 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 261 | parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') 262 | parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS') 263 | parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') 264 | parser.add_argument('--task', default='val', help="'val', 'test', 'study'") 265 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 266 | parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') 267 | parser.add_argument('--augment', action='store_true', help='augmented inference') 268 | parser.add_argument('--merge', action='store_true', help='use Merge NMS') 269 | parser.add_argument('--verbose', action='store_true', help='report mAP by class') 270 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 271 | opt = parser.parse_args() 272 | opt.save_json |= opt.data.endswith('coco.yaml') 273 | opt.data = check_file(opt.data) # check file 274 | print(opt) 275 | 276 | if opt.task in ['val', 'test']: # run normally 277 | test(opt.data, 278 | opt.weights, 279 | opt.batch_size, 280 | opt.max_box, 281 | opt.img_size, 282 | opt.conf_thres, 283 | opt.iou_thres, 284 | opt.save_json, 285 | opt.single_cls, 286 | opt.augment, 287 | opt.verbose) 288 | 289 | elif opt.task == 'study': # run over a range of settings and save/plot 290 | for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 291 | f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to 292 | x = list(range(320, 800, 64)) # x axis 293 | y = [] # y axis 294 | for i in x: # img-size 295 | print('\nRunning %s point %s...' % (f, i)) 296 | r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json) 297 | y.append(r + t) # results and times 298 | np.savetxt(f, y, fmt='%10.4g') # save 299 | os.system('zip -r study.zip study_*.txt') 300 | # utils.general.plot_study_txt(f, x) # plot 301 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import glob 3 | import logging 4 | import math 5 | import os 6 | import random 7 | import shutil 8 | import time 9 | from pathlib import Path 10 | 11 | import numpy as np 12 | import torch.distributed as dist 13 | import torch.nn.functional as F 14 | import torch.optim as optim 15 | import torch.optim.lr_scheduler as lr_scheduler 16 | import torch.utils.data 17 | import yaml 18 | from torch.cuda import amp 19 | from torch.nn.parallel import DistributedDataParallel as DDP 20 | from torch.utils.tensorboard import SummaryWriter 21 | from tqdm import tqdm 22 | 23 | import test # import test.py to get mAP after each epoch 24 | from models.yolo import Model 25 | from utils.datasets import create_dataloader 26 | from utils.general import ( 27 | torch_distributed_zero_first, labels_to_class_weights, plot_labels, check_anchors, labels_to_image_weights, 28 | compute_loss, plot_images, fitness, strip_optimizer, plot_results, get_latest_run, check_dataset, check_file, 29 | check_git_status, check_img_size, increment_dir, print_mutation, plot_evolution, set_logging) 30 | from utils.google_utils import attempt_download 31 | from utils.torch_utils import init_seeds, ModelEMA, select_device, intersect_dicts 32 | 33 | logger = logging.getLogger(__name__) 34 | 35 | 36 | def train(hyp, opt, device, tb_writer=None): 37 | logger.info(f'Hyperparameters {hyp}') 38 | log_dir = Path(tb_writer.log_dir) if tb_writer else Path(opt.logdir) / 'evolve' # logging directory 39 | wdir = log_dir / 'weights' # weights directory 40 | os.makedirs(wdir, exist_ok=True) 41 | last = wdir / 'last.pt' 42 | best = wdir / 'best.pt' 43 | results_file = str(log_dir / 'results.txt') 44 | epochs, batch_size, total_batch_size, weights, rank = \ 45 | opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank 46 | 47 | # Save run settings 48 | with open(log_dir / 'hyp.yaml', 'w') as f: 49 | yaml.dump(hyp, f, sort_keys=False) 50 | with open(log_dir / 'opt.yaml', 'w') as f: 51 | yaml.dump(vars(opt), f, sort_keys=False) 52 | 53 | # Configure 54 | cuda = device.type != 'cpu' 55 | init_seeds(2 + rank) 56 | with open(opt.data) as f: 57 | data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict 58 | with torch_distributed_zero_first(rank): 59 | check_dataset(data_dict) # check 60 | train_path = data_dict['train'] 61 | test_path = data_dict['val'] 62 | nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) # number classes, names 63 | assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check 64 | 65 | # Model 66 | pretrained = weights.endswith('.pt') 67 | if pretrained: 68 | with torch_distributed_zero_first(rank): 69 | attempt_download(weights) # download if not found locally 70 | ckpt = torch.load(weights, map_location=device) # load checkpoint 71 | if hyp.get('anchors'): 72 | ckpt['model'].yaml['anchors'] = round(hyp['anchors']) # force autoanchor 73 | model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device) # create 74 | exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [] # exclude keys 75 | state_dict = ckpt['model'].float().state_dict() # to FP32 76 | state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect 77 | model.load_state_dict(state_dict, strict=False) # load 78 | logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report 79 | else: 80 | model = Model(opt.cfg, ch=3, nc=nc).to(device) # create 81 | 82 | # Freeze 83 | freeze = ['', ] # parameter names to freeze (full or partial) 84 | if any(freeze): 85 | for k, v in model.named_parameters(): 86 | if any(x in k for x in freeze): 87 | print('freezing %s' % k) 88 | v.requires_grad = False 89 | 90 | # Optimizer 91 | nbs = 64 # nominal batch size 92 | accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing 93 | hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay 94 | 95 | pg0, pg1, pg2 = [], [], [] # optimizer parameter groups 96 | for k, v in model.named_parameters(): 97 | v.requires_grad = True 98 | if '.bias' in k: 99 | pg2.append(v) # biases 100 | elif '.weight' in k and '.bn' not in k: 101 | pg1.append(v) # apply weight decay 102 | else: 103 | pg0.append(v) # all else 104 | 105 | if opt.adam: 106 | optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum 107 | else: 108 | optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) 109 | 110 | optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay 111 | optimizer.add_param_group({'params': pg2}) # add pg2 (biases) 112 | logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0))) 113 | del pg0, pg1, pg2 114 | 115 | # Scheduler https://arxiv.org/pdf/1812.01187.pdf 116 | # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR 117 | lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf'] # cosine 118 | scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) 119 | # plot_lr_scheduler(optimizer, scheduler, epochs) 120 | 121 | # Resume 122 | start_epoch, best_fitness = 0, 0.0 123 | if pretrained: 124 | # Optimizer 125 | if ckpt['optimizer'] is not None: 126 | optimizer.load_state_dict(ckpt['optimizer']) 127 | best_fitness = ckpt['best_fitness'] 128 | 129 | # Results 130 | if ckpt.get('training_results') is not None: 131 | with open(results_file, 'w') as file: 132 | file.write(ckpt['training_results']) # write results.txt 133 | 134 | # Epochs 135 | start_epoch = ckpt['epoch'] + 1 136 | if opt.resume: 137 | assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs) 138 | shutil.copytree(wdir, wdir.parent / f'weights_backup_epoch{start_epoch - 1}') # save previous weights 139 | if epochs < start_epoch: 140 | logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' % 141 | (weights, ckpt['epoch'], epochs)) 142 | epochs += ckpt['epoch'] # finetune additional epochs 143 | 144 | del ckpt, state_dict 145 | 146 | # Image sizes 147 | gs = int(max(model.stride)) # grid size (max stride) 148 | imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples 149 | 150 | # DP mode 151 | if cuda and rank == -1 and torch.cuda.device_count() > 1: 152 | model = torch.nn.DataParallel(model) 153 | 154 | # SyncBatchNorm 155 | if opt.sync_bn and cuda and rank != -1: 156 | model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) 157 | logger.info('Using SyncBatchNorm()') 158 | 159 | # Exponential moving average 160 | ema = ModelEMA(model) if rank in [-1, 0] else None 161 | 162 | # DDP mode 163 | if cuda and rank != -1: 164 | model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank) 165 | 166 | # Trainloader 167 | dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, 168 | hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank, 169 | world_size=opt.world_size, workers=opt.workers) 170 | mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class 171 | nb = len(dataloader) # number of batches 172 | assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1) 173 | 174 | # Process 0 175 | if rank in [-1, 0]: 176 | ema.updates = start_epoch * nb // accumulate # set EMA updates 177 | testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, 178 | hyp=hyp, augment=False, cache=opt.cache_images, rect=True, rank=-1, 179 | world_size=opt.world_size, workers=opt.workers)[0] # testloader 180 | 181 | if not opt.resume: 182 | labels = np.concatenate(dataset.labels, 0) 183 | c = torch.tensor(labels[:, 0]) # classes 184 | # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency 185 | # model._initialize_biases(cf.to(device)) 186 | plot_labels(labels, save_dir=log_dir) 187 | if tb_writer: 188 | # tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384 189 | tb_writer.add_histogram('classes', c, 0) 190 | 191 | # Anchors 192 | if not opt.noautoanchor: 193 | check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) 194 | 195 | # Model parameters 196 | hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset 197 | model.nc = nc # attach number of classes to model 198 | model.hyp = hyp # attach hyperparameters to model 199 | model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou) 200 | model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights 201 | model.names = names 202 | 203 | # Start training 204 | t0 = time.time() 205 | nw = max(3 * nb, 1e3) # number of warmup iterations, max(3 epochs, 1k iterations) 206 | # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training 207 | maps = np.zeros(nc) # mAP per class 208 | results = (0, 0, 0, 0, 0, 0, 0) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification' 209 | scheduler.last_epoch = start_epoch - 1 # do not move 210 | scaler = amp.GradScaler(enabled=cuda) 211 | logger.info('Image sizes %g train, %g test\nUsing %g dataloader workers\nLogging results to %s\n' 212 | 'Starting training for %g epochs...' % (imgsz, imgsz_test, dataloader.num_workers, log_dir, epochs)) 213 | for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ 214 | model.train() 215 | 216 | # Update image weights (optional) 217 | if opt.image_weights: 218 | # Generate indices 219 | if rank in [-1, 0]: 220 | cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 # class weights 221 | iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights 222 | dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx 223 | # Broadcast if DDP 224 | if rank != -1: 225 | indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int() 226 | dist.broadcast(indices, 0) 227 | if rank != 0: 228 | dataset.indices = indices.cpu().numpy() 229 | 230 | # Update mosaic border 231 | # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) 232 | # dataset.mosaic_border = [b - imgsz, -b] # height, width borders 233 | 234 | mloss = torch.zeros(4, device=device) # mean losses 235 | if rank != -1: 236 | dataloader.sampler.set_epoch(epoch) 237 | pbar = enumerate(dataloader) 238 | logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls', 'total', 'targets', 'img_size')) 239 | if rank in [-1, 0]: 240 | pbar = tqdm(pbar, total=nb) # progress bar 241 | optimizer.zero_grad() 242 | for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- 243 | ni = i + nb * epoch # number integrated batches (since train start) 244 | imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0 245 | 246 | # Warmup 247 | if ni <= nw: 248 | xi = [0, nw] # x interp 249 | # model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou) 250 | accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round()) 251 | for j, x in enumerate(optimizer.param_groups): 252 | # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 253 | x['lr'] = np.interp(ni, xi, [0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)]) 254 | if 'momentum' in x: 255 | x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']]) 256 | 257 | # Multi-scale 258 | if opt.multi_scale: 259 | sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size 260 | sf = sz / max(imgs.shape[2:]) # scale factor 261 | if sf != 1: 262 | ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) 263 | imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) 264 | 265 | # Forward 266 | with amp.autocast(enabled=cuda): 267 | pred = model(imgs) # forward 268 | loss, loss_items = compute_loss(pred, targets.to(device), model) # loss scaled by batch_size 269 | if rank != -1: 270 | loss *= opt.world_size # gradient averaged between devices in DDP mode 271 | 272 | # Backward 273 | scaler.scale(loss).backward() 274 | 275 | # Optimize 276 | if ni % accumulate == 0: 277 | scaler.step(optimizer) # optimizer.step 278 | scaler.update() 279 | optimizer.zero_grad() 280 | if ema: 281 | ema.update(model) 282 | 283 | # Print 284 | if rank in [-1, 0]: 285 | mloss = (mloss * i + loss_items) / (i + 1) # update mean losses 286 | mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB) 287 | s = ('%10s' * 2 + '%10.4g' * 6) % ( 288 | '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1]) 289 | pbar.set_description(s) 290 | 291 | # Plot 292 | if ni < 3: 293 | f = str(log_dir / ('train_batch%g.jpg' % ni)) # filename 294 | result = plot_images(images=imgs, targets=targets, paths=paths, fname=f) 295 | if tb_writer and result is not None: 296 | tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) 297 | # tb_writer.add_graph(model, imgs) # add model to tensorboard 298 | 299 | # end batch ------------------------------------------------------------------------------------------------ 300 | 301 | # Scheduler 302 | lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard 303 | scheduler.step() 304 | 305 | # DDP process 0 or single-GPU 306 | if rank in [-1, 0]: 307 | # mAP 308 | if ema: 309 | ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride']) 310 | final_epoch = epoch + 1 == epochs 311 | if not opt.notest or final_epoch: # Calculate mAP 312 | if final_epoch: # replot predictions 313 | [os.remove(x) for x in glob.glob(str(log_dir / 'test_batch*_pred.jpg')) if os.path.exists(x)] 314 | results, maps, times = test.test(opt.data, 315 | batch_size=total_batch_size, 316 | imgsz=imgsz_test, 317 | model=ema.ema, 318 | single_cls=opt.single_cls, 319 | dataloader=testloader, 320 | save_dir=log_dir) 321 | 322 | # Write 323 | with open(results_file, 'a') as f: 324 | f.write(s + '%10.4g' * 7 % results + '\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls) 325 | if len(opt.name) and opt.bucket: 326 | os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name)) 327 | 328 | # Tensorboard 329 | if tb_writer: 330 | tags = ['train/giou_loss', 'train/obj_loss', 'train/cls_loss', # train loss 331 | 'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 332 | 'val/giou_loss', 'val/obj_loss', 'val/cls_loss', # val loss 333 | 'x/lr0', 'x/lr1', 'x/lr2'] # params 334 | for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags): 335 | tb_writer.add_scalar(tag, x, epoch) 336 | 337 | # Update best mAP 338 | fi = fitness(np.array(results).reshape(1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1] 339 | if fi > best_fitness: 340 | best_fitness = fi 341 | 342 | # Save model 343 | save = (not opt.nosave) or (final_epoch and not opt.evolve) 344 | if save: 345 | with open(results_file, 'r') as f: # create checkpoint 346 | ckpt = {'epoch': epoch, 347 | 'best_fitness': best_fitness, 348 | 'training_results': f.read(), 349 | 'model': ema.ema, 350 | 'optimizer': None if final_epoch else optimizer.state_dict()} 351 | 352 | # Save last, best and delete 353 | torch.save(ckpt, last) 354 | if best_fitness == fi: 355 | torch.save(ckpt, best) 356 | del ckpt 357 | # end epoch ---------------------------------------------------------------------------------------------------- 358 | # end training 359 | 360 | if rank in [-1, 0]: 361 | # Strip optimizers 362 | n = opt.name if opt.name.isnumeric() else '' 363 | fresults, flast, fbest = log_dir / f'results{n}.txt', wdir / f'last{n}.pt', wdir / f'best{n}.pt' 364 | for f1, f2 in zip([wdir / 'last.pt', wdir / 'best.pt', results_file], [flast, fbest, fresults]): 365 | if os.path.exists(f1): 366 | os.rename(f1, f2) # rename 367 | if str(f2).endswith('.pt'): # is *.pt 368 | strip_optimizer(f2) # strip optimizer 369 | os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket else None # upload 370 | # Finish 371 | if not opt.evolve: 372 | plot_results(save_dir=log_dir) # save as results.png 373 | logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600)) 374 | 375 | dist.destroy_process_group() if rank not in [-1, 0] else None 376 | torch.cuda.empty_cache() 377 | return results 378 | 379 | 380 | if __name__ == '__main__': 381 | parser = argparse.ArgumentParser() 382 | parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path') 383 | parser.add_argument('--cfg', type=str, default='', help='model.yaml path') 384 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path') 385 | parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path') 386 | parser.add_argument('--epochs', type=int, default=300) 387 | parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs') 388 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') 389 | parser.add_argument('--rect', action='store_true', help='rectangular training') 390 | parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') 391 | parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') 392 | parser.add_argument('--notest', action='store_true', help='only test final epoch') 393 | parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check') 394 | parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') 395 | parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') 396 | parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') 397 | parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') 398 | parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied') 399 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 400 | parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') 401 | parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') 402 | parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') 403 | parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') 404 | parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') 405 | parser.add_argument('--logdir', type=str, default='runs/', help='logging directory') 406 | parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers') 407 | opt = parser.parse_args() 408 | 409 | # Set DDP variables 410 | opt.total_batch_size = opt.batch_size 411 | opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1 412 | opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1 413 | set_logging(opt.global_rank) 414 | if opt.global_rank in [-1, 0]: 415 | check_git_status() 416 | 417 | # Resume 418 | if opt.resume: # resume an interrupted run 419 | ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path 420 | log_dir = Path(ckpt).parent.parent # runs/exp0 421 | assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' 422 | with open(log_dir / 'opt.yaml') as f: 423 | opt = argparse.Namespace(**yaml.load(f, Loader=yaml.FullLoader)) # replace 424 | opt.cfg, opt.weights, opt.resume = '', ckpt, True 425 | logger.info('Resuming training from %s' % ckpt) 426 | 427 | else: 428 | # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml') 429 | opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp) # check files 430 | assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified' 431 | opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test) 432 | log_dir = increment_dir(Path(opt.logdir) / 'exp', opt.name) # runs/exp1 433 | 434 | device = select_device(opt.device, batch_size=opt.batch_size) 435 | 436 | # DDP mode 437 | if opt.local_rank != -1: 438 | assert torch.cuda.device_count() > opt.local_rank 439 | torch.cuda.set_device(opt.local_rank) 440 | device = torch.device('cuda', opt.local_rank) 441 | dist.init_process_group(backend='nccl', init_method='env://') # distributed backend 442 | assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count' 443 | opt.batch_size = opt.total_batch_size // opt.world_size 444 | 445 | logger.info(opt) 446 | with open(opt.hyp) as f: 447 | hyp = yaml.load(f, Loader=yaml.FullLoader) # load hyps 448 | 449 | # Train 450 | if not opt.evolve: 451 | tb_writer = None 452 | if opt.global_rank in [-1, 0]: 453 | logger.info('Start Tensorboard with "tensorboard --logdir %s", view at http://localhost:6006/' % opt.logdir) 454 | tb_writer = SummaryWriter(log_dir=log_dir) # runs/exp0 455 | 456 | train(hyp, opt, device, tb_writer) 457 | 458 | # Evolve hyperparameters (optional) 459 | else: 460 | # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) 461 | meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) 462 | 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) 463 | 'momentum': (0.1, 0.6, 0.98), # SGD momentum/Adam beta1 464 | 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay 465 | 'giou': (1, 0.02, 0.2), # GIoU loss gain 466 | 'cls': (1, 0.2, 4.0), # cls loss gain 467 | 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight 468 | 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) 469 | 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight 470 | 'iou_t': (0, 0.1, 0.7), # IoU training threshold 471 | 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold 472 | 'anchors': (1, 2.0, 10.0), # anchors per output grid (0 to ignore) 473 | 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) 474 | 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) 475 | 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) 476 | 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) 477 | 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) 478 | 'translate': (1, 0.0, 0.9), # image translation (+/- fraction) 479 | 'scale': (1, 0.0, 0.9), # image scale (+/- gain) 480 | 'shear': (1, 0.0, 10.0), # image shear (+/- deg) 481 | 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 482 | 'flipud': (1, 0.0, 1.0), # image flip up-down (probability) 483 | 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) 484 | 'mixup': (1, 0.0, 1.0)} # image mixup (probability) 485 | 486 | assert opt.local_rank == -1, 'DDP mode not implemented for --evolve' 487 | opt.notest, opt.nosave = True, True # only test/save final epoch 488 | # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices 489 | yaml_file = Path('runs/evolve/hyp_evolved.yaml') # save best result here 490 | if opt.bucket: 491 | os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket) # download evolve.txt if exists 492 | 493 | for _ in range(1): # generations to evolve 494 | if os.path.exists('evolve.txt'): # if evolve.txt exists: select best hyps and mutate 495 | # Select parent(s) 496 | parent = 'single' # parent selection method: 'single' or 'weighted' 497 | x = np.loadtxt('evolve.txt', ndmin=2) 498 | n = min(5, len(x)) # number of previous results to consider 499 | x = x[np.argsort(-fitness(x))][:n] # top n mutations 500 | w = fitness(x) - fitness(x).min() # weights 501 | if parent == 'single' or len(x) == 1: 502 | # x = x[random.randint(0, n - 1)] # random selection 503 | x = x[random.choices(range(n), weights=w)[0]] # weighted selection 504 | elif parent == 'weighted': 505 | x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination 506 | 507 | # Mutate 508 | mp, s = 0.9, 0.2 # mutation probability, sigma 509 | npr = np.random 510 | npr.seed(int(time.time())) 511 | g = np.array([x[0] for x in meta.values()]) # gains 0-1 512 | ng = len(meta) 513 | v = np.ones(ng) 514 | while all(v == 1): # mutate until a change occurs (prevent duplicates) 515 | v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0) 516 | for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300) 517 | hyp[k] = float(x[i + 7] * v[i]) # mutate 518 | 519 | # Constrain to limits 520 | for k, v in meta.items(): 521 | hyp[k] = max(hyp[k], v[1]) # lower limit 522 | hyp[k] = min(hyp[k], v[2]) # upper limit 523 | hyp[k] = round(hyp[k], 5) # significant digits 524 | 525 | # Train mutation 526 | results = train(hyp.copy(), opt, device) 527 | 528 | # Write mutation results 529 | print_mutation(hyp.copy(), results, yaml_file, opt.bucket) 530 | 531 | # Plot results 532 | plot_evolution(yaml_file) 533 | print('Hyperparameter evolution complete. Best results saved as: %s\nCommand to train a new model with these ' 534 | 'hyperparameters: $ python train.py --hyp %s' % (yaml_file, yaml_file)) 535 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zzh-tju/yolov5/fca5b7b7a85bcadac394f99c1805e54d9c20a21f/utils/__init__.py -------------------------------------------------------------------------------- /utils/activations.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | # Swish https://arxiv.org/pdf/1905.02244.pdf --------------------------------------------------------------------------- 7 | class Swish(nn.Module): # 8 | @staticmethod 9 | def forward(x): 10 | return x * torch.sigmoid(x) 11 | 12 | 13 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish() 14 | @staticmethod 15 | def forward(x): 16 | # return x * F.hardsigmoid(x) # for torchscript and CoreML 17 | return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX 18 | 19 | 20 | class MemoryEfficientSwish(nn.Module): 21 | class F(torch.autograd.Function): 22 | @staticmethod 23 | def forward(ctx, x): 24 | ctx.save_for_backward(x) 25 | return x * torch.sigmoid(x) 26 | 27 | @staticmethod 28 | def backward(ctx, grad_output): 29 | x = ctx.saved_tensors[0] 30 | sx = torch.sigmoid(x) 31 | return grad_output * (sx * (1 + x * (1 - sx))) 32 | 33 | def forward(self, x): 34 | return self.F.apply(x) 35 | 36 | 37 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 38 | class Mish(nn.Module): 39 | @staticmethod 40 | def forward(x): 41 | return x * F.softplus(x).tanh() 42 | 43 | 44 | class MemoryEfficientMish(nn.Module): 45 | class F(torch.autograd.Function): 46 | @staticmethod 47 | def forward(ctx, x): 48 | ctx.save_for_backward(x) 49 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 50 | 51 | @staticmethod 52 | def backward(ctx, grad_output): 53 | x = ctx.saved_tensors[0] 54 | sx = torch.sigmoid(x) 55 | fx = F.softplus(x).tanh() 56 | return grad_output * (fx + x * sx * (1 - fx * fx)) 57 | 58 | def forward(self, x): 59 | return self.F.apply(x) 60 | 61 | 62 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 63 | class FReLU(nn.Module): 64 | def __init__(self, c1, k=3): # ch_in, kernel 65 | super().__init__() 66 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1) 67 | self.bn = nn.BatchNorm2d(c1) 68 | 69 | def forward(self, x): 70 | return torch.max(x, self.bn(self.conv(x))) 71 | -------------------------------------------------------------------------------- /utils/evolve.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Hyperparameter evolution commands (avoids CUDA memory leakage issues) 3 | # Replaces train.py python generations 'for' loop with a bash 'for' loop 4 | 5 | # Start on 4-GPU machine 6 | #for i in 0 1 2 3; do 7 | # t=ultralytics/yolov5:evolve && sudo docker pull $t && sudo docker run -d --ipc=host --gpus all -v "$(pwd)"/VOC:/usr/src/VOC $t bash utils/evolve.sh $i 8 | # sleep 60 # avoid simultaneous evolve.txt read/write 9 | #done 10 | 11 | # Hyperparameter evolution commands 12 | while true; do 13 | # python train.py --batch 64 --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50 --evolve --bucket ult/evolve/voc --device $1 14 | python train.py --batch 40 --weights yolov5m.pt --data coco.yaml --img 640 --epochs 30 --evolve --bucket ult/evolve/coco --device $1 15 | done 16 | -------------------------------------------------------------------------------- /utils/google_utils.py: -------------------------------------------------------------------------------- 1 | # This file contains google utils: https://cloud.google.com/storage/docs/reference/libraries 2 | # pip install --upgrade google-cloud-storage 3 | # from google.cloud import storage 4 | 5 | import os 6 | import platform 7 | import subprocess 8 | import time 9 | from pathlib import Path 10 | 11 | import torch 12 | 13 | 14 | def gsutil_getsize(url=''): 15 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 16 | s = subprocess.check_output('gsutil du %s' % url, shell=True).decode('utf-8') 17 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 18 | 19 | 20 | def attempt_download(weights): 21 | # Attempt to download pretrained weights if not found locally 22 | weights = weights.strip().replace("'", '') 23 | file = Path(weights).name 24 | 25 | msg = weights + ' missing, try downloading from https://github.com/ultralytics/yolov5/releases/' 26 | models = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'] # available models 27 | 28 | if file in models and not os.path.isfile(weights): 29 | # Google Drive 30 | # d = {'yolov5s.pt': '1R5T6rIyy3lLwgFXNms8whc-387H0tMQO', 31 | # 'yolov5m.pt': '1vobuEExpWQVpXExsJ2w-Mbf3HJjWkQJr', 32 | # 'yolov5l.pt': '1hrlqD1Wdei7UT4OgT785BEk1JwnSvNEV', 33 | # 'yolov5x.pt': '1mM8aZJlWTxOg7BZJvNUMrTnA2AbeCVzS'} 34 | # r = gdrive_download(id=d[file], name=weights) if file in d else 1 35 | # if r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6: # check 36 | # return 37 | 38 | try: # GitHub 39 | url = 'https://github.com/ultralytics/yolov5/releases/download/v3.0/' + file 40 | print('Downloading %s to %s...' % (url, weights)) 41 | torch.hub.download_url_to_file(url, weights) 42 | assert os.path.exists(weights) and os.path.getsize(weights) > 1E6 # check 43 | except Exception as e: # GCP 44 | print('Download error: %s' % e) 45 | url = 'https://storage.googleapis.com/ultralytics/yolov5/ckpt/' + file 46 | print('Downloading %s to %s...' % (url, weights)) 47 | r = os.system('curl -L %s -o %s' % (url, weights)) # torch.hub.download_url_to_file(url, weights) 48 | finally: 49 | if not (os.path.exists(weights) and os.path.getsize(weights) > 1E6): # check 50 | os.remove(weights) if os.path.exists(weights) else None # remove partial downloads 51 | print('ERROR: Download failure: %s' % msg) 52 | print('') 53 | return 54 | 55 | 56 | def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'): 57 | # Downloads a file from Google Drive. from utils.google_utils import *; gdrive_download() 58 | t = time.time() 59 | 60 | print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='') 61 | os.remove(name) if os.path.exists(name) else None # remove existing 62 | os.remove('cookie') if os.path.exists('cookie') else None 63 | 64 | # Attempt file download 65 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 66 | os.system('curl -c ./cookie -s -L "drive.google.com/uc?export=download&id=%s" > %s ' % (id, out)) 67 | if os.path.exists('cookie'): # large file 68 | s = 'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm=%s&id=%s" -o %s' % (get_token(), id, name) 69 | else: # small file 70 | s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id) 71 | r = os.system(s) # execute, capture return 72 | os.remove('cookie') if os.path.exists('cookie') else None 73 | 74 | # Error check 75 | if r != 0: 76 | os.remove(name) if os.path.exists(name) else None # remove partial 77 | print('Download error ') # raise Exception('Download error') 78 | return r 79 | 80 | # Unzip if archive 81 | if name.endswith('.zip'): 82 | print('unzipping... ', end='') 83 | os.system('unzip -q %s' % name) # unzip 84 | os.remove(name) # remove zip to free space 85 | 86 | print('Done (%.1fs)' % (time.time() - t)) 87 | return r 88 | 89 | 90 | def get_token(cookie="./cookie"): 91 | with open(cookie) as f: 92 | for line in f: 93 | if "download" in line: 94 | return line.split()[-1] 95 | return "" 96 | 97 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 98 | # # Uploads a file to a bucket 99 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 100 | # 101 | # storage_client = storage.Client() 102 | # bucket = storage_client.get_bucket(bucket_name) 103 | # blob = bucket.blob(destination_blob_name) 104 | # 105 | # blob.upload_from_filename(source_file_name) 106 | # 107 | # print('File {} uploaded to {}.'.format( 108 | # source_file_name, 109 | # destination_blob_name)) 110 | # 111 | # 112 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 113 | # # Uploads a blob from a bucket 114 | # storage_client = storage.Client() 115 | # bucket = storage_client.get_bucket(bucket_name) 116 | # blob = bucket.blob(source_blob_name) 117 | # 118 | # blob.download_to_filename(destination_file_name) 119 | # 120 | # print('Blob {} downloaded to {}.'.format( 121 | # source_blob_name, 122 | # destination_file_name)) 123 | -------------------------------------------------------------------------------- /utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import math 3 | import os 4 | import time 5 | from copy import deepcopy 6 | 7 | import torch 8 | import torch.backends.cudnn as cudnn 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import torchvision.models as models 12 | 13 | logger = logging.getLogger(__name__) 14 | 15 | 16 | def init_seeds(seed=0): 17 | torch.manual_seed(seed) 18 | 19 | # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html 20 | if seed == 0: # slower, more reproducible 21 | cudnn.deterministic = True 22 | cudnn.benchmark = False 23 | else: # faster, less reproducible 24 | cudnn.deterministic = False 25 | cudnn.benchmark = True 26 | 27 | 28 | def select_device(device='', batch_size=None): 29 | # device = 'cpu' or '0' or '0,1,2,3' 30 | cpu_request = device.lower() == 'cpu' 31 | if device and not cpu_request: # if device requested other than 'cpu' 32 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable 33 | assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity 34 | 35 | cuda = False if cpu_request else torch.cuda.is_available() 36 | if cuda: 37 | c = 1024 ** 2 # bytes to MB 38 | ng = torch.cuda.device_count() 39 | if ng > 1 and batch_size: # check that batch_size is compatible with device_count 40 | assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng) 41 | x = [torch.cuda.get_device_properties(i) for i in range(ng)] 42 | s = 'Using CUDA ' 43 | for i in range(0, ng): 44 | if i == 1: 45 | s = ' ' * len(s) 46 | logger.info("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % 47 | (s, i, x[i].name, x[i].total_memory / c)) 48 | else: 49 | logger.info('Using CPU') 50 | 51 | logger.info('') # skip a line 52 | return torch.device('cuda:0' if cuda else 'cpu') 53 | 54 | 55 | def time_synchronized(): 56 | torch.cuda.synchronize() if torch.cuda.is_available() else None 57 | return time.time() 58 | 59 | 60 | def is_parallel(model): 61 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 62 | 63 | 64 | def intersect_dicts(da, db, exclude=()): 65 | # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values 66 | return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape} 67 | 68 | 69 | def initialize_weights(model): 70 | for m in model.modules(): 71 | t = type(m) 72 | if t is nn.Conv2d: 73 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 74 | elif t is nn.BatchNorm2d: 75 | m.eps = 1e-3 76 | m.momentum = 0.03 77 | elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 78 | m.inplace = True 79 | 80 | 81 | def find_modules(model, mclass=nn.Conv2d): 82 | # Finds layer indices matching module class 'mclass' 83 | return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)] 84 | 85 | 86 | def sparsity(model): 87 | # Return global model sparsity 88 | a, b = 0., 0. 89 | for p in model.parameters(): 90 | a += p.numel() 91 | b += (p == 0).sum() 92 | return b / a 93 | 94 | 95 | def prune(model, amount=0.3): 96 | # Prune model to requested global sparsity 97 | import torch.nn.utils.prune as prune 98 | print('Pruning model... ', end='') 99 | for name, m in model.named_modules(): 100 | if isinstance(m, nn.Conv2d): 101 | prune.l1_unstructured(m, name='weight', amount=amount) # prune 102 | prune.remove(m, 'weight') # make permanent 103 | print(' %.3g global sparsity' % sparsity(model)) 104 | 105 | 106 | def fuse_conv_and_bn(conv, bn): 107 | # https://tehnokv.com/posts/fusing-batchnorm-and-conv/ 108 | with torch.no_grad(): 109 | # init 110 | fusedconv = nn.Conv2d(conv.in_channels, 111 | conv.out_channels, 112 | kernel_size=conv.kernel_size, 113 | stride=conv.stride, 114 | padding=conv.padding, 115 | groups=conv.groups, 116 | bias=True).to(conv.weight.device) 117 | 118 | # prepare filters 119 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 120 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 121 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) 122 | 123 | # prepare spatial bias 124 | b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias 125 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) 126 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) 127 | 128 | return fusedconv 129 | 130 | 131 | def model_info(model, verbose=False): 132 | # Plots a line-by-line description of a PyTorch model 133 | n_p = sum(x.numel() for x in model.parameters()) # number parameters 134 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients 135 | if verbose: 136 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) 137 | for i, (name, p) in enumerate(model.named_parameters()): 138 | name = name.replace('module_list.', '') 139 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' % 140 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) 141 | 142 | try: # FLOPS 143 | from thop import profile 144 | flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, 64, 64),), verbose=False)[0] / 1E9 * 2 145 | fs = ', %.1f GFLOPS' % (flops * 100) # 640x640 FLOPS 146 | except: 147 | fs = '' 148 | 149 | logger.info( 150 | 'Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs)) 151 | 152 | 153 | def load_classifier(name='resnet101', n=2): 154 | # Loads a pretrained model reshaped to n-class output 155 | model = models.__dict__[name](pretrained=True) 156 | 157 | # Display model properties 158 | input_size = [3, 224, 224] 159 | input_space = 'RGB' 160 | input_range = [0, 1] 161 | mean = [0.485, 0.456, 0.406] 162 | std = [0.229, 0.224, 0.225] 163 | for x in ['input_size', 'input_space', 'input_range', 'mean', 'std']: 164 | print(x + ' =', eval(x)) 165 | 166 | # Reshape output to n classes 167 | filters = model.fc.weight.shape[1] 168 | model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True) 169 | model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True) 170 | model.fc.out_features = n 171 | return model 172 | 173 | 174 | def scale_img(img, ratio=1.0, same_shape=False): # img(16,3,256,416), r=ratio 175 | # scales img(bs,3,y,x) by ratio 176 | if ratio == 1.0: 177 | return img 178 | else: 179 | h, w = img.shape[2:] 180 | s = (int(h * ratio), int(w * ratio)) # new size 181 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize 182 | if not same_shape: # pad/crop img 183 | gs = 32 # (pixels) grid size 184 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] 185 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean 186 | 187 | 188 | def copy_attr(a, b, include=(), exclude=()): 189 | # Copy attributes from b to a, options to only include [...] and to exclude [...] 190 | for k, v in b.__dict__.items(): 191 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 192 | continue 193 | else: 194 | setattr(a, k, v) 195 | 196 | 197 | class ModelEMA: 198 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 199 | Keep a moving average of everything in the model state_dict (parameters and buffers). 200 | This is intended to allow functionality like 201 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 202 | A smoothed version of the weights is necessary for some training schemes to perform well. 203 | This class is sensitive where it is initialized in the sequence of model init, 204 | GPU assignment and distributed training wrappers. 205 | """ 206 | 207 | def __init__(self, model, decay=0.9999, updates=0): 208 | # Create EMA 209 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA 210 | # if next(model.parameters()).device.type != 'cpu': 211 | # self.ema.half() # FP16 EMA 212 | self.updates = updates # number of EMA updates 213 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) 214 | for p in self.ema.parameters(): 215 | p.requires_grad_(False) 216 | 217 | def update(self, model): 218 | # Update EMA parameters 219 | with torch.no_grad(): 220 | self.updates += 1 221 | d = self.decay(self.updates) 222 | 223 | msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict 224 | for k, v in self.ema.state_dict().items(): 225 | if v.dtype.is_floating_point: 226 | v *= d 227 | v += (1. - d) * msd[k].detach() 228 | 229 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): 230 | # Update EMA attributes 231 | copy_attr(self.ema, model, include, exclude) 232 | -------------------------------------------------------------------------------- /weights/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Download common models 3 | 4 | python -c " 5 | from utils.google_utils import *; 6 | attempt_download('weights/yolov5s.pt'); 7 | attempt_download('weights/yolov5m.pt'); 8 | attempt_download('weights/yolov5l.pt'); 9 | attempt_download('weights/yolov5x.pt') 10 | " 11 | --------------------------------------------------------------------------------