├── models ├── __init__.py ├── experimental.py └── yolo.py ├── data ├── oxfordhand.names ├── valid.cache ├── oxfordhand.data ├── coco.data ├── coco_128img.data ├── get_coco_dataset_gdrive.sh ├── coco.names ├── get_coco_dataset.sh └── coco_128img.txt ├── utils ├── __init__.py ├── parse_config.py ├── google_utils.py ├── torch_utils.py ├── downloads.py ├── autoanchor.py ├── gcp.sh ├── adabound.py └── metrics.py ├── tk1_time.xls ├── prune_yolov5s.sh ├── shortcut_prune_yolov5s.sh ├── slim_prune_yolov5s.sh ├── slim_prune_yolov5s_8x.sh ├── README.md ├── .gitignore ├── prune_yolov5s.py ├── cfg └── yolov5s_v6_hand.cfg ├── LICENSE ├── slim_prune_yolov5s.py ├── shortcut_prune_yolov5s.py ├── slim_prune_yolov5s_8x.py ├── test.py └── test_yolov5s.py /models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/oxfordhand.names: -------------------------------------------------------------------------------- 1 | hand -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | # -------------------------------------------------------------------------------- /tk1_time.xls: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZJU-lishuang/yolov5_prune/HEAD/tk1_time.xls -------------------------------------------------------------------------------- /data/valid.cache: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ZJU-lishuang/yolov5_prune/HEAD/data/valid.cache -------------------------------------------------------------------------------- /data/oxfordhand.data: -------------------------------------------------------------------------------- 1 | classes= 1 2 | train=data/train.txt 3 | valid=data/valid.txt 4 | names=data/oxfordhand.names 5 | -------------------------------------------------------------------------------- /data/coco.data: -------------------------------------------------------------------------------- 1 | classes=80 2 | train=../coco/trainvalno5k.txt 3 | valid=../coco/5k.txt 4 | names=data/coco.names 5 | backup=backup/ 6 | eval=coco 7 | -------------------------------------------------------------------------------- /prune_yolov5s.sh: -------------------------------------------------------------------------------- 1 | python prune_yolov5s.py --cfg cfg/yolov5s_v6_hand.cfg --data data/oxfordhand.data --weights weights/last_v6s_0.pt --percent 0.8 --img_size 640 -------------------------------------------------------------------------------- /data/coco_128img.data: -------------------------------------------------------------------------------- 1 | classes=80 2 | train=./data/coco_128img.txt 3 | valid=./data/coco_128img.txt 4 | names=data/coco.names 5 | backup=backup/ 6 | eval=coco 7 | -------------------------------------------------------------------------------- /shortcut_prune_yolov5s.sh: -------------------------------------------------------------------------------- 1 | python shortcut_prune_yolov5s.py --cfg cfg/yolov5s_v6_hand.cfg --data data/oxfordhand.data --weights weights/last_v6s.pt --percent 0.5 --img_size 640 -------------------------------------------------------------------------------- /slim_prune_yolov5s.sh: -------------------------------------------------------------------------------- 1 | python slim_prune_yolov5s.py --cfg cfg/yolov5s_v6_hand.cfg --data data/oxfordhand.data --weights weights/last_v6s.pt --global_percent 0.5 --layer_keep 0.01 --img_size 640 -------------------------------------------------------------------------------- /slim_prune_yolov5s_8x.sh: -------------------------------------------------------------------------------- 1 | python slim_prune_yolov5s_8x.py --cfg cfg/yolov5s_v6_hand.cfg --data data/oxfordhand.data --weights weights/last_v6s.pt --global_percent 0.5 --layer_keep 0.01 --img_size 640 -------------------------------------------------------------------------------- /data/get_coco_dataset_gdrive.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # https://stackoverflow.com/questions/48133080/how-to-download-a-google-drive-url-via-curl-or-wget/48133859 3 | 4 | # Zip coco folder 5 | # zip -r coco.zip coco 6 | # tar -czvf coco.tar.gz coco 7 | 8 | # Set fileid and filename 9 | filename="coco.zip" 10 | fileid="1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO" # coco.zip 11 | 12 | # Download from Google Drive, accepting presented query 13 | curl -c ./cookie -s -L "https://drive.google.com/uc?export=download&id=${fileid}" > /dev/null 14 | curl -Lb ./cookie "https://drive.google.com/uc?export=download&confirm=`awk '/download/ {print $NF}' ./cookie`&id=${fileid}" -o ${filename} 15 | rm ./cookie 16 | 17 | # Unzip 18 | unzip -q ${filename} # for coco.zip 19 | # tar -xzf ${filename} # for coco.tar.gz 20 | -------------------------------------------------------------------------------- /data/coco.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorcycle 5 | airplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | couch 59 | potted plant 60 | bed 61 | dining table 62 | toilet 63 | tv 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush 81 | -------------------------------------------------------------------------------- /utils/parse_config.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | def parse_model_cfg(path): 5 | # Parses the yolo-v3 layer configuration file and returns module definitions 6 | file = open(path, 'r') 7 | lines = file.read().split('\n') 8 | lines = [x for x in lines if x and not x.startswith('#')] 9 | lines = [x.rstrip().lstrip() for x in lines] # get rid of fringe whitespaces 10 | mdefs = [] # module definitions 11 | for line in lines: 12 | if line.startswith('['): # This marks the start of a new block 13 | mdefs.append({}) 14 | mdefs[-1]['type'] = line[1:-1].rstrip() 15 | if mdefs[-1]['type'] == 'convolutional': 16 | mdefs[-1]['batch_normalize'] = 0 # pre-populate with zeros (may be overwritten later) 17 | else: 18 | key, val = line.split("=") 19 | key = key.rstrip() 20 | 21 | if 'anchors' in key: 22 | mdefs[-1][key] = np.array([float(x) for x in val.split(',')]).reshape((-1, 2)) # np anchors 23 | else: 24 | mdefs[-1][key] = val.strip() 25 | 26 | return mdefs 27 | 28 | 29 | def parse_data_cfg(path): 30 | # Parses the data configuration file 31 | options = dict() 32 | with open(path, 'r') as fp: 33 | lines = fp.readlines() 34 | 35 | for line in lines: 36 | line = line.strip() 37 | if line == '' or line.startswith('#'): 38 | continue 39 | key, val = line.split('=') 40 | options[key.strip()] = val.strip() 41 | 42 | return options 43 | 44 | -------------------------------------------------------------------------------- /data/get_coco_dataset.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # CREDIT: https://github.com/pjreddie/darknet/tree/master/scripts/get_coco_dataset.sh 3 | 4 | # Clone COCO API 5 | git clone https://github.com/pdollar/coco && cd coco 6 | 7 | # Download Images 8 | mkdir images && cd images 9 | wget -c https://pjreddie.com/media/files/train2014.zip 10 | wget -c https://pjreddie.com/media/files/val2014.zip 11 | 12 | # Unzip 13 | unzip -q train2014.zip 14 | unzip -q val2014.zip 15 | 16 | # (optional) Delete zip files 17 | rm -rf *.zip 18 | 19 | cd .. 20 | 21 | # Download COCO Metadata 22 | wget -c https://pjreddie.com/media/files/instances_train-val2014.zip 23 | wget -c https://pjreddie.com/media/files/coco/5k.part 24 | wget -c https://pjreddie.com/media/files/coco/trainvalno5k.part 25 | wget -c https://pjreddie.com/media/files/coco/labels.tgz 26 | tar xzf labels.tgz 27 | unzip -q instances_train-val2014.zip 28 | 29 | # Set Up Image Lists 30 | paste <(awk "{print \"$PWD\"}" <5k.part) 5k.part | tr -d '\t' > 5k.txt 31 | paste <(awk "{print \"$PWD\"}" trainvalno5k.txt 32 | 33 | # get xview training data 34 | # wget -O train_images.tgz 'https://d307kc0mrhucc3.cloudfront.net/train_images.tgz?Expires=1530124049&Signature=JrQoxipmsETvb7eQHCfDFUO-QEHJGAayUv0i-ParmS-1hn7hl9D~bzGuHWG82imEbZSLUARTtm0wOJ7EmYMGmG5PtLKz9H5qi6DjoSUuFc13NQ-~6yUhE~NfPaTnehUdUMCa3On2wl1h1ZtRG~0Jq1P-AJbpe~oQxbyBrs1KccaMa7FK4F4oMM6sMnNgoXx8-3O77kYw~uOpTMFmTaQdHln6EztW0Lx17i57kK3ogbSUpXgaUTqjHCRA1dWIl7PY1ngQnLslkLhZqmKcaL-BvWf0ZGjHxCDQBpnUjIlvMu5NasegkwD9Jjc0ClgTxsttSkmbapVqaVC8peR0pO619Q__&Key-Pair-Id=APKAIKGDJB5C3XUL2DXQ' 35 | # tar -xvzf train_images.tgz 36 | # sudo rm -rf train_images/._* 37 | # lastly convert each .tif to a .bmp for faster loading in cv2 38 | 39 | # ./coco/images/train2014/COCO_train2014_000000167126.jpg # corrupted image 40 | -------------------------------------------------------------------------------- /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 | 4 | import os 5 | import time 6 | 7 | 8 | # from google.cloud import storage 9 | 10 | 11 | def gdrive_download(id='1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO', name='coco.zip'): 12 | # https://gist.github.com/tanaikech/f0f2d122e05bf5f971611258c22c110f 13 | # Downloads a file from Google Drive, accepting presented query 14 | # from utils.google_utils import *; gdrive_download() 15 | t = time.time() 16 | 17 | print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='') 18 | if os.path.exists(name): # remove existing 19 | os.remove(name) 20 | 21 | # Attempt large file download 22 | s = ["curl -c ./cookie -s -L \"https://drive.google.com/uc?export=download&id=%s\" > /dev/null" % id, 23 | "curl -Lb ./cookie -s \"https://drive.google.com/uc?export=download&confirm=`awk '/download/ {print $NF}' ./cookie`&id=%s\" -o %s" % ( 24 | id, name), 25 | 'rm ./cookie'] 26 | [os.system(x) for x in s] # run commands 27 | 28 | # Attempt small file download 29 | if not os.path.exists(name): # file size < 40MB 30 | s = 'curl -f -L -o %s https://drive.google.com/uc?export=download&id=%s' % (name, id) 31 | os.system(s) 32 | 33 | # Unzip if archive 34 | if name.endswith('.zip'): 35 | print('unzipping... ', end='') 36 | os.system('unzip -q %s' % name) # unzip 37 | os.remove(name) # remove zip to free space 38 | 39 | print('Done (%.1fs)' % (time.time() - t)) 40 | 41 | 42 | def upload_blob(bucket_name, source_file_name, destination_blob_name): 43 | # Uploads a file to a bucket 44 | # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 45 | 46 | storage_client = storage.Client() 47 | bucket = storage_client.get_bucket(bucket_name) 48 | blob = bucket.blob(destination_blob_name) 49 | 50 | blob.upload_from_filename(source_file_name) 51 | 52 | print('File {} uploaded to {}.'.format( 53 | source_file_name, 54 | destination_blob_name)) 55 | 56 | 57 | def download_blob(bucket_name, source_blob_name, destination_file_name): 58 | # Uploads a blob from a bucket 59 | storage_client = storage.Client() 60 | bucket = storage_client.get_bucket(bucket_name) 61 | blob = bucket.blob(source_blob_name) 62 | 63 | blob.download_to_filename(destination_file_name) 64 | 65 | print('Blob {} downloaded to {}.'.format( 66 | source_blob_name, 67 | destination_file_name)) 68 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # yolov5_prune 2 | 本项目基于[tanluren/yolov3-channel-and-layer-pruning](https://github.com/tanluren/yolov3-channel-and-layer-pruning)实现，将项目扩展到yolov5上。
3 | 4 | 项目的基本流程是，使用[ultralytics/yolov5](https://github.com/ultralytics/yolov5)训练自己的数据集，在模型性能达到要求但速度未达到要求时，对模型进行剪枝。首先是稀疏化训练，稀疏化训练很重要，如果模型稀疏度不够，剪枝比例过大会导致剪枝后的模型map接近0。剪枝完成后对模型进行微调回复精度。
5 | 6 | 本项目使用的yolov5为第六版本。
7 | yolov5第三版本参考[yolov5-v4-prune](https://github.com/ZJU-lishuang/yolov5_prune/tree/v4)
8 | yolov5第三版本参考[yolov5-v3-prune](https://github.com/ZJU-lishuang/yolov5_prune/tree/v3)
9 | yolov5第二版本参考[yolov5-v2-prune](https://github.com/ZJU-lishuang/yolov5_prune/tree/v2)
10 | 11 | PS：在开源数据集和自有数据集上模型均剪枝成功。 12 | 13 | ## 实例流程 14 | 数据集下载[dataset](http://www.robots.ox.ac.uk/~vgg/data/hands/downloads/hand_dataset.tar.gz)
15 | 数据集转为可训练格式[converter](https://github.com/ZJU-lishuang/yolov5-v4/blob/main/data/converter.py) 16 | ### STEP1:基础训练 17 | 附件：[训练记录](https://drive.google.com/drive/folders/1ZdgYUk5B9-KsE8m-CyhFv0-jzURm2SCV?usp=sharing)
18 | ### STEP2:稀疏训练 19 | 附件：[稀疏训练记录](https://drive.google.com/drive/folders/1-aUNG_spznsF-KJ9nsur4r7XtZds4rU0?usp=sharing)
20 | ### STEP3:八倍通道剪枝 21 | 附件：[剪枝后模型](https://drive.google.com/drive/folders/1KJYsVlaB5_3QZB3r0nzJUKYW_oTHW4Pa?usp=sharing)
22 | ### STEP4:微调finetune 23 | 附件：[微调训练记录](https://drive.google.com/drive/folders/1AsHG_w--NdSPCV4sPaPYpcOnMyOpNgHx?usp=sharing)
24 | ### STEP4:微调finetune，使用蒸馏技术优化模型，效果优于单纯的微调模型 25 | 附件：[微调蒸馏训练记录](https://drive.google.com/drive/folders/1VDVHwhPReIN5WNLeb-8wnGmZbpe7pc_c?usp=sharing)
26 | 27 | ## 剪枝步骤 28 | #### STEP1:基础训练 29 | **项目**[yolov5](https://github.com/ZJU-lishuang/yolov5-v6)
30 | 示例代码
31 | ``` 32 | python train.py --img 640 --batch 16 --epochs 50 --weights weights/yolov5s_v6.pt --data data/coco_hand.yaml --cfg models/yolov5s.yaml --name s_hand 33 | ``` 34 | 35 | #### STEP2:稀疏训练 36 | --prune 0 适用于通道剪枝策略一，--prune 1 适用于其他剪枝策略。
37 | **项目**[yolov5](https://github.com/ZJU-lishuang/yolov5-v6)
38 | 示例代码
39 | ``` 40 | python train_sparsity.py --img 640 --batch 16 --epochs 50 --data data/coco_hand.yaml --cfg models/yolov5s.yaml --weights runs/train/s_hand/weights/last.pt --name s_hand_sparsity -sr --scale 0.001 --prune 1 41 | ``` 42 | 43 | #### STEP3:通道剪枝策略一 44 | 不对shortcut直连的层进行剪枝，避免维度处理。
45 | ``` 46 | python prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/oxfordhand.data --weights weights/yolov5s_prune0.pt --percent 0.8 47 | ``` 48 | 49 | #### STEP3:通道剪枝策略二 50 | 对shortcut层也进行了剪枝，剪枝采用每组shortcut中第一个卷积层的mask。
51 | ``` 52 | python shortcut_prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/oxfordhand.data --weights weights/yolov5s_prune1.pt --percent 0.3 53 | ``` 54 | 55 | #### STEP3:通道剪枝策略三 56 | 先以全局阈值找出各卷积层的mask，然后对于每组shortcut，它将相连的各卷积层的剪枝mask取并集，用merge后的mask进行剪枝。
57 | ``` 58 | python slim_prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/oxfordhand.data --weights weights/yolov5s_prune1.pt --global_percent 0.8 --layer_keep 0.01 59 | ``` 60 | 61 | #### STEP3:八倍通道剪枝 62 | 在硬件部署上发现，模型剪枝率相同时，通道数为8的倍数速度最快。（采坑：需要将硬件性能开启到最大）
63 | 示例代码
64 | ``` 65 | python slim_prune_yolov5s_8x.py --cfg cfg/yolov5s_v6_hand.cfg --data data/oxfordhand.data --weights weights/last_v6s.pt --global_percent 0.5 --layer_keep 0.01 --img_size 640 66 | ``` 67 | 68 | #### STEP4:微调finetune 69 | **项目**[yolov5](https://github.com/ZJU-lishuang/yolov5-v6)
70 | 示例代码
71 | ``` 72 | python prune_finetune.py --img 640 --batch 16 --epochs 50 --data data/coco_hand.yaml --cfg ./cfg/prune_0.6_keep_0.01_8x_yolov5s_v6_hand.cfg --weights ./weights/prune_0.6_keep_0.01_8x_last_v6s.pt --name s_hand_finetune 73 | ``` 74 | 75 | #### STEP4:微调finetune，使用蒸馏技术优化模型 76 | **项目**[yolov5](https://github.com/ZJU-lishuang/yolov5-v6)
77 | 示例代码
78 | ``` 79 | python prune_finetune.py --img 640 --batch 16 --epochs 50 --data data/coco_hand.yaml --cfg ./cfg/prune_0.6_keep_0.01_8x_yolov5s_v6_hand.cfg --weights ./weights/prune_0.6_keep_0.01_8x_last_v6s.pt --name s_hand_finetune_distill --distill 80 | ``` 81 | 82 | #### STEP5:剪枝后模型推理 83 | **项目**[yolov5](https://github.com/ZJU-lishuang/yolov5-v6)
84 | 示例代码
85 | ```shell 86 | python prune_detect.py --weights weights/last_s_hand_finetune.pt --img 640 --conf 0.7 --source inference/images 87 | ``` 88 | 89 | 90 | -------------------------------------------------------------------------------- /.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 | *.shapes 21 | *.json 22 | 23 | #data 24 | data/coco128 25 | data/hand_dataset 26 | 27 | cfg/prune_*.cfg 28 | 29 | 30 | # *.cfg 31 | !cfg/yolov3*.cfg 32 | 33 | storage.googleapis.com 34 | runs/* 35 | 36 | 37 | pycocotools/* 38 | results*.txt 39 | gcp_test*.sh 40 | 41 | # MATLAB GitIgnore ----------------------------------------------------------------------------------------------------- 42 | *.m~ 43 | *.mat 44 | !targets*.mat 45 | 46 | # Neural Network weights ----------------------------------------------------------------------------------------------- 47 | *.weights 48 | *.pt 49 | *.onnx 50 | *.mlmodel 51 | *.torchscript 52 | darknet53.conv.74 53 | yolov3-tiny.conv.15 54 | 55 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 56 | # Byte-compiled / optimized / DLL files 57 | __pycache__/ 58 | *.py[cod] 59 | *$py.class 60 | 61 | # C extensions 62 | *.so 63 | 64 | # Distribution / packaging 65 | .Python 66 | env/ 67 | build/ 68 | develop-eggs/ 69 | dist/ 70 | downloads/ 71 | eggs/ 72 | .eggs/ 73 | lib/ 74 | lib64/ 75 | parts/ 76 | sdist/ 77 | var/ 78 | wheels/ 79 | *.egg-info/ 80 | .installed.cfg 81 | *.egg 82 | 83 | # PyInstaller 84 | # Usually these files are written by a python script from a template 85 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 86 | *.manifest 87 | *.spec 88 | 89 | # Installer logs 90 | pip-log.txt 91 | pip-delete-this-directory.txt 92 | 93 | # Unit test / coverage reports 94 | htmlcov/ 95 | .tox/ 96 | .coverage 97 | .coverage.* 98 | .cache 99 | nosetests.xml 100 | coverage.xml 101 | *.cover 102 | .hypothesis/ 103 | 104 | # Translations 105 | *.mo 106 | *.pot 107 | 108 | # Django stuff: 109 | *.log 110 | local_settings.py 111 | 112 | # Flask stuff: 113 | instance/ 114 | .webassets-cache 115 | 116 | # Scrapy stuff: 117 | .scrapy 118 | 119 | # Sphinx documentation 120 | docs/_build/ 121 | 122 | # PyBuilder 123 | target/ 124 | 125 | # Jupyter Notebook 126 | .ipynb_checkpoints 127 | 128 | # pyenv 129 | .python-version 130 | 131 | # celery beat schedule file 132 | celerybeat-schedule 133 | 134 | # SageMath parsed files 135 | *.sage.py 136 | 137 | # dotenv 138 | .env 139 | 140 | # virtualenv 141 | .venv 142 | venv/ 143 | ENV/ 144 | 145 | # Spyder project settings 146 | .spyderproject 147 | .spyproject 148 | 149 | # Rope project settings 150 | .ropeproject 151 | 152 | # mkdocs documentation 153 | /site 154 | 155 | # mypy 156 | .mypy_cache/ 157 | 158 | 159 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 160 | 161 | # General 162 | .DS_Store 163 | .AppleDouble 164 | .LSOverride 165 | 166 | # Icon must end with two \r 167 | Icon 168 | Icon? 169 | 170 | # Thumbnails 171 | ._* 172 | 173 | # Files that might appear in the root of a volume 174 | .DocumentRevisions-V100 175 | .fseventsd 176 | .Spotlight-V100 177 | .TemporaryItems 178 | .Trashes 179 | .VolumeIcon.icns 180 | .com.apple.timemachine.donotpresent 181 | 182 | # Directories potentially created on remote AFP share 183 | .AppleDB 184 | .AppleDesktop 185 | Network Trash Folder 186 | Temporary Items 187 | .apdisk 188 | 189 | 190 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 191 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 192 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 193 | 194 | # User-specific stuff: 195 | .idea/* 196 | .idea/**/workspace.xml 197 | .idea/**/tasks.xml 198 | .idea/dictionaries 199 | .html # Bokeh Plots 200 | .pg # TensorFlow Frozen Graphs 201 | .avi # videos 202 | 203 | # Sensitive or high-churn files: 204 | .idea/**/dataSources/ 205 | .idea/**/dataSources.ids 206 | .idea/**/dataSources.local.xml 207 | .idea/**/sqlDataSources.xml 208 | .idea/**/dynamic.xml 209 | .idea/**/uiDesigner.xml 210 | 211 | # Gradle: 212 | .idea/**/gradle.xml 213 | .idea/**/libraries 214 | 215 | # CMake 216 | cmake-build-debug/ 217 | cmake-build-release/ 218 | 219 | # Mongo Explorer plugin: 220 | .idea/**/mongoSettings.xml 221 | 222 | ## File-based project format: 223 | *.iws 224 | 225 | ## Plugin-specific files: 226 | 227 | # IntelliJ 228 | out/ 229 | 230 | # mpeltonen/sbt-idea plugin 231 | .idea_modules/ 232 | 233 | # JIRA plugin 234 | atlassian-ide-plugin.xml 235 | 236 | # Cursive Clojure plugin 237 | .idea/replstate.xml 238 | 239 | # Crashlytics plugin (for Android Studio and IntelliJ) 240 | com_crashlytics_export_strings.xml 241 | crashlytics.properties 242 | crashlytics-build.properties 243 | fabric.properties 244 | 245 | tensorboard/ 246 | -------------------------------------------------------------------------------- /models/experimental.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Experimental modules 4 | """ 5 | 6 | import numpy as np 7 | import torch 8 | import torch.nn as nn 9 | 10 | from models.common import Conv 11 | from utils.downloads import attempt_download 12 | 13 | 14 | class CrossConv(nn.Module): 15 | # Cross Convolution Downsample 16 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 17 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 18 | super().__init__() 19 | c_ = int(c2 * e) # hidden channels 20 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 21 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 22 | self.add = shortcut and c1 == c2 23 | 24 | def forward(self, x): 25 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 26 | 27 | 28 | class Sum(nn.Module): 29 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 30 | def __init__(self, n, weight=False): # n: number of inputs 31 | super().__init__() 32 | self.weight = weight # apply weights boolean 33 | self.iter = range(n - 1) # iter object 34 | if weight: 35 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights 36 | 37 | def forward(self, x): 38 | y = x[0] # no weight 39 | if self.weight: 40 | w = torch.sigmoid(self.w) * 2 41 | for i in self.iter: 42 | y = y + x[i + 1] * w[i] 43 | else: 44 | for i in self.iter: 45 | y = y + x[i + 1] 46 | return y 47 | 48 | 49 | class MixConv2d(nn.Module): 50 | # Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595 51 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): 52 | super().__init__() 53 | groups = len(k) 54 | if equal_ch: # equal c_ per group 55 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices 56 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels 57 | else: # equal weight.numel() per group 58 | b = [c2] + [0] * groups 59 | a = np.eye(groups + 1, groups, k=-1) 60 | a -= np.roll(a, 1, axis=1) 61 | a *= np.array(k) ** 2 62 | a[0] = 1 63 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 64 | 65 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) 66 | self.bn = nn.BatchNorm2d(c2) 67 | self.act = nn.LeakyReLU(0.1, inplace=True) 68 | 69 | def forward(self, x): 70 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 71 | 72 | 73 | class Ensemble(nn.ModuleList): 74 | # Ensemble of models 75 | def __init__(self): 76 | super().__init__() 77 | 78 | def forward(self, x, augment=False, profile=False, visualize=False): 79 | y = [] 80 | for module in self: 81 | y.append(module(x, augment, profile, visualize)[0]) 82 | # y = torch.stack(y).max(0)[0] # max ensemble 83 | # y = torch.stack(y).mean(0) # mean ensemble 84 | y = torch.cat(y, 1) # nms ensemble 85 | return y, None # inference, train output 86 | 87 | 88 | def attempt_load(weights, map_location=None, inplace=True, fuse=True): 89 | from models.yolo import Detect, Model 90 | 91 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 92 | model = Ensemble() 93 | for w in weights if isinstance(weights, list) else [weights]: 94 | ckpt = torch.load(attempt_download(w), map_location=map_location) # load 95 | if fuse: 96 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model 97 | else: 98 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval()) # without layer fuse 99 | 100 | 101 | # Compatibility updates 102 | for m in model.modules(): 103 | if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model]: 104 | m.inplace = inplace # pytorch 1.7.0 compatibility 105 | if type(m) is Detect: 106 | if not isinstance(m.anchor_grid, list): # new Detect Layer compatibility 107 | delattr(m, 'anchor_grid') 108 | setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl) 109 | elif type(m) is Conv: 110 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 111 | 112 | if len(model) == 1: 113 | return model[-1] # return model 114 | else: 115 | print(f'Ensemble created with {weights}\n') 116 | for k in ['names']: 117 | setattr(model, k, getattr(model[-1], k)) 118 | model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride 119 | return model # return ensemble 120 | -------------------------------------------------------------------------------- /utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416) 8 | # scales img(bs,3,y,x) by ratio constrained to gs-multiple 9 | if ratio == 1.0: 10 | return img 11 | else: 12 | h, w = img.shape[2:] 13 | s = (int(h * ratio), int(w * ratio)) # new size 14 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize 15 | if not same_shape: # pad/crop img 16 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] 17 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean 18 | 19 | def copy_attr(a, b, include=(), exclude=()): 20 | # Copy attributes from b to a, options to only include [...] and to exclude [...] 21 | for k, v in b.__dict__.items(): 22 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 23 | continue 24 | else: 25 | setattr(a, k, v) 26 | 27 | def time_sync(): 28 | # pytorch-accurate time 29 | if torch.cuda.is_available(): 30 | torch.cuda.synchronize() 31 | return time.time() 32 | 33 | def init_seeds(seed=0): 34 | torch.manual_seed(seed) 35 | torch.cuda.manual_seed(seed) 36 | torch.cuda.manual_seed_all(seed) 37 | 38 | # Remove randomness (may be slower on Tesla GPUs) # https://pytorch.org/docs/stable/notes/randomness.html 39 | if seed == 0: 40 | torch.backends.cudnn.deterministic = True 41 | torch.backends.cudnn.benchmark = False 42 | 43 | 44 | def select_device(device='', apex=False): 45 | # device = 'cpu' or '0' or '0,1,2,3' 46 | cpu_request = device.lower() == 'cpu' 47 | if device and not cpu_request: # if device requested other than 'cpu' 48 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable 49 | assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity 50 | 51 | cuda = False if cpu_request else torch.cuda.is_available() 52 | if cuda: 53 | c = 1024 ** 2 # bytes to MB 54 | ng = torch.cuda.device_count() 55 | x = [torch.cuda.get_device_properties(i) for i in range(ng)] 56 | cuda_str = 'Using CUDA ' + ('Apex ' if apex else '') # apex for mixed precision https://github.com/NVIDIA/apex 57 | for i in range(0, ng): 58 | if i == 1: 59 | cuda_str = ' ' * len(cuda_str) 60 | print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % 61 | (cuda_str, i, x[i].name, x[i].total_memory / c)) 62 | else: 63 | print('Using CPU') 64 | 65 | print('') # skip a line 66 | return torch.device('cuda:0' if cuda else 'cpu') 67 | 68 | def time_synchronized(): 69 | torch.cuda.synchronize() if torch.cuda.is_available() else None 70 | return time.time() 71 | 72 | 73 | def is_parallel(model): 74 | # is model is parallel with DP or DDP 75 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 76 | 77 | def initialize_weights(model): 78 | for m in model.modules(): 79 | t = type(m) 80 | if t is nn.Conv2d: 81 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 82 | elif t is nn.BatchNorm2d: 83 | m.eps = 1e-3 84 | m.momentum = 0.03 85 | elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 86 | m.inplace = True 87 | 88 | def fuse_conv_and_bn(conv, bn): 89 | # https://tehnokv.com/posts/fusing-batchnorm-and-conv/ 90 | with torch.no_grad(): 91 | # init 92 | fusedconv = torch.nn.Conv2d(conv.in_channels, 93 | conv.out_channels, 94 | kernel_size=conv.kernel_size, 95 | stride=conv.stride, 96 | padding=conv.padding, 97 | bias=True) 98 | 99 | # prepare filters 100 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 101 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 102 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) 103 | 104 | # prepare spatial bias 105 | if conv.bias is not None: 106 | b_conv = conv.bias 107 | else: 108 | b_conv = torch.zeros(conv.weight.size(0)) 109 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) 110 | fusedconv.bias.copy_(b_conv + b_bn) 111 | 112 | return fusedconv 113 | 114 | 115 | def model_info(model, report='summary'): 116 | # Plots a line-by-line description of a PyTorch model 117 | n_p = sum(x.numel() for x in model.parameters()) # number parameters 118 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients 119 | if report is 'full': 120 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) 121 | for i, (name, p) in enumerate(model.named_parameters()): 122 | name = name.replace('module_list.', '') 123 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' % 124 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) 125 | print('Model Summary: %g layers, %g parameters, %g gradients' % (len(list(model.parameters())), n_p, n_g)) 126 | -------------------------------------------------------------------------------- /utils/downloads.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Download utils 4 | """ 5 | 6 | import os 7 | import platform 8 | import subprocess 9 | import time 10 | import urllib 11 | from pathlib import Path 12 | from zipfile import ZipFile 13 | 14 | import requests 15 | import torch 16 | 17 | 18 | def gsutil_getsize(url=''): 19 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 20 | s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8') 21 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 22 | 23 | 24 | def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''): 25 | # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes 26 | file = Path(file) 27 | assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}" 28 | try: # url1 29 | print(f'Downloading {url} to {file}...') 30 | torch.hub.download_url_to_file(url, str(file)) 31 | assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check 32 | except Exception as e: # url2 33 | file.unlink(missing_ok=True) # remove partial downloads 34 | print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...') 35 | os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail 36 | finally: 37 | if not file.exists() or file.stat().st_size < min_bytes: # check 38 | file.unlink(missing_ok=True) # remove partial downloads 39 | print(f"ERROR: {assert_msg}\n{error_msg}") 40 | print('') 41 | 42 | 43 | def attempt_download(file, repo='ultralytics/yolov5'): # from utils.downloads import *; attempt_download() 44 | # Attempt file download if does not exist 45 | file = Path(str(file).strip().replace("'", '')) 46 | 47 | if not file.exists(): 48 | # URL specified 49 | name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc. 50 | if str(file).startswith(('http:/', 'https:/')): # download 51 | url = str(file).replace(':/', '://') # Pathlib turns :// -> :/ 52 | name = name.split('?')[0] # parse authentication https://url.com/file.txt?auth... 53 | safe_download(file=name, url=url, min_bytes=1E5) 54 | return name 55 | 56 | # GitHub assets 57 | file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required) 58 | try: 59 | response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api 60 | assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] 61 | tag = response['tag_name'] # i.e. 'v1.0' 62 | except: # fallback plan 63 | assets = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 64 | 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt'] 65 | try: 66 | tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1] 67 | except: 68 | tag = 'v5.0' # current release 69 | 70 | if name in assets: 71 | safe_download(file, 72 | url=f'https://github.com/{repo}/releases/download/{tag}/{name}', 73 | # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional) 74 | min_bytes=1E5, 75 | error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/') 76 | 77 | return str(file) 78 | 79 | 80 | def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'): 81 | # Downloads a file from Google Drive. from yolov5.utils.downloads import *; gdrive_download() 82 | t = time.time() 83 | file = Path(file) 84 | cookie = Path('cookie') # gdrive cookie 85 | print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='') 86 | file.unlink(missing_ok=True) # remove existing file 87 | cookie.unlink(missing_ok=True) # remove existing cookie 88 | 89 | # Attempt file download 90 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 91 | os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}') 92 | if os.path.exists('cookie'): # large file 93 | s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}' 94 | else: # small file 95 | s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"' 96 | r = os.system(s) # execute, capture return 97 | cookie.unlink(missing_ok=True) # remove existing cookie 98 | 99 | # Error check 100 | if r != 0: 101 | file.unlink(missing_ok=True) # remove partial 102 | print('Download error ') # raise Exception('Download error') 103 | return r 104 | 105 | # Unzip if archive 106 | if file.suffix == '.zip': 107 | print('unzipping... ', end='') 108 | ZipFile(file).extractall(path=file.parent) # unzip 109 | file.unlink() # remove zip 110 | 111 | print(f'Done ({time.time() - t:.1f}s)') 112 | return r 113 | 114 | 115 | def get_token(cookie="./cookie"): 116 | with open(cookie) as f: 117 | for line in f: 118 | if "download" in line: 119 | return line.split()[-1] 120 | return "" 121 | 122 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries ---------------------------------------------- 123 | # 124 | # 125 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 126 | # # Uploads a file to a bucket 127 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 128 | # 129 | # storage_client = storage.Client() 130 | # bucket = storage_client.get_bucket(bucket_name) 131 | # blob = bucket.blob(destination_blob_name) 132 | # 133 | # blob.upload_from_filename(source_file_name) 134 | # 135 | # print('File {} uploaded to {}.'.format( 136 | # source_file_name, 137 | # destination_blob_name)) 138 | # 139 | # 140 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 141 | # # Uploads a blob from a bucket 142 | # storage_client = storage.Client() 143 | # bucket = storage_client.get_bucket(bucket_name) 144 | # blob = bucket.blob(source_blob_name) 145 | # 146 | # blob.download_to_filename(destination_file_name) 147 | # 148 | # print('Blob {} downloaded to {}.'.format( 149 | # source_blob_name, 150 | # destination_file_name)) 151 | -------------------------------------------------------------------------------- /utils/autoanchor.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Auto-anchor utils 4 | """ 5 | 6 | import random 7 | 8 | import numpy as np 9 | import torch 10 | import yaml 11 | from tqdm import tqdm 12 | 13 | from utils.general import colorstr 14 | 15 | 16 | def check_anchor_order(m): 17 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary 18 | a = m.anchors.prod(-1).view(-1) # anchor area 19 | da = a[-1] - a[0] # delta a 20 | ds = m.stride[-1] - m.stride[0] # delta s 21 | if da.sign() != ds.sign(): # same order 22 | print('Reversing anchor order') 23 | m.anchors[:] = m.anchors.flip(0) 24 | 25 | 26 | def check_anchors(dataset, model, thr=4.0, imgsz=640): 27 | # Check anchor fit to data, recompute if necessary 28 | prefix = colorstr('autoanchor: ') 29 | print(f'\n{prefix}Analyzing anchors... ', end='') 30 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() 31 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) 32 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale 33 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh 34 | 35 | def metric(k): # compute metric 36 | r = wh[:, None] / k[None] 37 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 38 | best = x.max(1)[0] # best_x 39 | aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold 40 | bpr = (best > 1. / thr).float().mean() # best possible recall 41 | return bpr, aat 42 | 43 | anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors 44 | bpr, aat = metric(anchors.cpu().view(-1, 2)) 45 | print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='') 46 | if bpr < 0.98: # threshold to recompute 47 | print('. Attempting to improve anchors, please wait...') 48 | na = m.anchors.numel() // 2 # number of anchors 49 | try: 50 | anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) 51 | except Exception as e: 52 | print(f'{prefix}ERROR: {e}') 53 | new_bpr = metric(anchors)[0] 54 | if new_bpr > bpr: # replace anchors 55 | anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors) 56 | m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss 57 | check_anchor_order(m) 58 | print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.') 59 | else: 60 | print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.') 61 | print('') # newline 62 | 63 | 64 | def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): 65 | """ Creates kmeans-evolved anchors from training dataset 66 | 67 | Arguments: 68 | dataset: path to data.yaml, or a loaded dataset 69 | n: number of anchors 70 | img_size: image size used for training 71 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 72 | gen: generations to evolve anchors using genetic algorithm 73 | verbose: print all results 74 | 75 | Return: 76 | k: kmeans evolved anchors 77 | 78 | Usage: 79 | from utils.autoanchor import *; _ = kmean_anchors() 80 | """ 81 | from scipy.cluster.vq import kmeans 82 | 83 | thr = 1. / thr 84 | prefix = colorstr('autoanchor: ') 85 | 86 | def metric(k, wh): # compute metrics 87 | r = wh[:, None] / k[None] 88 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 89 | # x = wh_iou(wh, torch.tensor(k)) # iou metric 90 | return x, x.max(1)[0] # x, best_x 91 | 92 | def anchor_fitness(k): # mutation fitness 93 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh) 94 | return (best * (best > thr).float()).mean() # fitness 95 | 96 | def print_results(k): 97 | k = k[np.argsort(k.prod(1))] # sort small to large 98 | x, best = metric(k, wh0) 99 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr 100 | print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr') 101 | print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' 102 | f'past_thr={x[x > thr].mean():.3f}-mean: ', end='') 103 | for i, x in enumerate(k): 104 | print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg 105 | return k 106 | 107 | if isinstance(dataset, str): # *.yaml file 108 | with open(dataset, errors='ignore') as f: 109 | data_dict = yaml.safe_load(f) # model dict 110 | from utils.datasets import LoadImagesAndLabels 111 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) 112 | 113 | # Get label wh 114 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) 115 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh 116 | 117 | # Filter 118 | i = (wh0 < 3.0).any(1).sum() 119 | if i: 120 | print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.') 121 | wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels 122 | # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1 123 | 124 | # Kmeans calculation 125 | print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...') 126 | s = wh.std(0) # sigmas for whitening 127 | k, dist = kmeans(wh / s, n, iter=30) # points, mean distance 128 | assert len(k) == n, f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}' 129 | k *= s 130 | wh = torch.tensor(wh, dtype=torch.float32) # filtered 131 | wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered 132 | k = print_results(k) 133 | 134 | # Plot 135 | # k, d = [None] * 20, [None] * 20 136 | # for i in tqdm(range(1, 21)): 137 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance 138 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) 139 | # ax = ax.ravel() 140 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') 141 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh 142 | # ax[0].hist(wh[wh[:, 0]<100, 0],400) 143 | # ax[1].hist(wh[wh[:, 1]<100, 1],400) 144 | # fig.savefig('wh.png', dpi=200) 145 | 146 | # Evolve 147 | npr = np.random 148 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma 149 | pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar 150 | for _ in pbar: 151 | v = np.ones(sh) 152 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates) 153 | v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) 154 | kg = (k.copy() * v).clip(min=2.0) 155 | fg = anchor_fitness(kg) 156 | if fg > f: 157 | f, k = fg, kg.copy() 158 | pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}' 159 | if verbose: 160 | print_results(k) 161 | 162 | return print_results(k) 163 | -------------------------------------------------------------------------------- /utils/gcp.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # New VM 4 | rm -rf sample_data yolov3 darknet apex coco cocoapi knife knifec 5 | git clone https://github.com/ultralytics/yolov3 6 | # git clone https://github.com/AlexeyAB/darknet && cd darknet && make GPU=1 CUDNN=1 CUDNN_HALF=1 OPENCV=0 && wget -c https://pjreddie.com/media/files/darknet53.conv.74 && cd .. 7 | git clone https://github.com/NVIDIA/apex && cd apex && pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" . --user && cd .. && rm -rf apex 8 | # git clone https://github.com/cocodataset/cocoapi && cd cocoapi/PythonAPI && make && cd ../.. && cp -r cocoapi/PythonAPI/pycocotools yolov3 9 | sudo conda install -y -c conda-forge scikit-image tensorboard pycocotools 10 | python3 -c " 11 | from yolov3.utils.google_utils import gdrive_download 12 | gdrive_download('1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO','coco.zip')" 13 | sudo shutdown 14 | 15 | # Re-clone 16 | rm -rf yolov3 # Warning: remove existing 17 | git clone https://github.com/ultralytics/yolov3 && cd yolov3 # master 18 | # git clone -b test --depth 1 https://github.com/ultralytics/yolov3 test # branch 19 | python3 train.py --img-size 320 --weights weights/darknet53.conv.74 --epochs 27 --batch-size 64 --accumulate 1 20 | 21 | # Train 22 | python3 train.py 23 | 24 | # Resume 25 | python3 train.py --resume 26 | 27 | # Detect 28 | python3 detect.py 29 | 30 | # Test 31 | python3 test.py --save-json 32 | 33 | # Evolve 34 | for i in {0..500} 35 | do 36 | python3 train.py --data data/coco.data --img-size 320 --epochs 1 --batch-size 64 --accumulate 1 --evolve --bucket yolov4 37 | done 38 | 39 | # Git pull 40 | git pull https://github.com/ultralytics/yolov3 # master 41 | git pull https://github.com/ultralytics/yolov3 test # branch 42 | 43 | # Test Darknet training 44 | python3 test.py --weights ../darknet/backup/yolov3.backup 45 | 46 | # Copy last.pt TO bucket 47 | gsutil cp yolov3/weights/last1gpu.pt gs://ultralytics 48 | 49 | # Copy last.pt FROM bucket 50 | gsutil cp gs://ultralytics/last.pt yolov3/weights/last.pt 51 | wget https://storage.googleapis.com/ultralytics/yolov3/last_v1_0.pt -O weights/last_v1_0.pt 52 | wget https://storage.googleapis.com/ultralytics/yolov3/best_v1_0.pt -O weights/best_v1_0.pt 53 | 54 | # Reproduce tutorials 55 | rm results*.txt # WARNING: removes existing results 56 | python3 train.py --nosave --data data/coco_1img.data && mv results.txt results0r_1img.txt 57 | python3 train.py --nosave --data data/coco_10img.data && mv results.txt results0r_10img.txt 58 | python3 train.py --nosave --data data/coco_100img.data && mv results.txt results0r_100img.txt 59 | # python3 train.py --nosave --data data/coco_100img.data --transfer && mv results.txt results3_100imgTL.txt 60 | python3 -c "from utils import utils; utils.plot_results()" 61 | # gsutil cp results*.txt gs://ultralytics 62 | gsutil cp results.png gs://ultralytics 63 | sudo shutdown 64 | 65 | # Reproduce mAP 66 | python3 test.py --save-json --img-size 608 67 | python3 test.py --save-json --img-size 416 68 | python3 test.py --save-json --img-size 320 69 | sudo shutdown 70 | 71 | # Benchmark script 72 | git clone https://github.com/ultralytics/yolov3 # clone our repo 73 | git clone https://github.com/NVIDIA/apex && cd apex && pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" . --user && cd .. && rm -rf apex # install nvidia apex 74 | python3 -c "from yolov3.utils.google_utils import gdrive_download; gdrive_download('1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO','coco.zip')" # download coco dataset (20GB) 75 | cd yolov3 && clear && python3 train.py --epochs 1 # run benchmark (~30 min) 76 | 77 | # Unit tests 78 | python3 detect.py # detect 2 persons, 1 tie 79 | python3 test.py --data data/coco_32img.data # test mAP = 0.8 80 | python3 train.py --data data/coco_32img.data --epochs 5 --nosave # train 5 epochs 81 | python3 train.py --data data/coco_1cls.data --epochs 5 --nosave # train 5 epochs 82 | python3 train.py --data data/coco_1img.data --epochs 5 --nosave # train 5 epochs 83 | 84 | # AlexyAB Darknet 85 | gsutil cp -r gs://sm6/supermarket2 . # dataset from bucket 86 | rm -rf darknet && git clone https://github.com/AlexeyAB/darknet && cd darknet && wget -c https://pjreddie.com/media/files/darknet53.conv.74 # sudo apt install libopencv-dev && make 87 | ./darknet detector calc_anchors data/coco_img64.data -num_of_clusters 9 -width 320 -height 320 # kmeans anchor calculation 88 | ./darknet detector train ../supermarket2/supermarket2.data ../yolo_v3_spp_pan_scale.cfg darknet53.conv.74 -map -dont_show # train spp 89 | ./darknet detector train ../yolov3/data/coco.data ../yolov3-spp.cfg darknet53.conv.74 -map -dont_show # train spp coco 90 | 91 | ./darknet detector train data/coco.data ../yolov3-spp.cfg darknet53.conv.74 -map -dont_show # train spp 92 | gsutil cp -r backup/*5000.weights gs://sm6/weights 93 | sudo shutdown 94 | 95 | 96 | ./darknet detector train ../supermarket2/supermarket2.data ../yolov3-tiny-sm2-1cls.cfg yolov3-tiny.conv.15 -map -dont_show # train tiny 97 | ./darknet detector train ../supermarket2/supermarket2.data cfg/yolov3-spp-sm2-1cls.cfg backup/yolov3-spp-sm2-1cls_last.weights # resume 98 | python3 train.py --data ../supermarket2/supermarket2.data --cfg ../yolov3-spp-sm2-1cls.cfg --epochs 100 --num-workers 8 --img-size 320 --nosave # train ultralytics 99 | python3 test.py --data ../supermarket2/supermarket2.data --weights ../darknet/backup/yolov3-spp-sm2-1cls_5000.weights --cfg cfg/yolov3-spp-sm2-1cls.cfg # test 100 | gsutil cp -r backup/*.weights gs://sm6/weights # weights to bucket 101 | 102 | python3 test.py --data ../supermarket2/supermarket2.data --weights weights/yolov3-spp-sm2-1cls_5000.weights --cfg ../yolov3-spp-sm2-1cls.cfg --img-size 320 --conf-thres 0.2 # test 103 | python3 test.py --data ../supermarket2/supermarket2.data --weights weights/yolov3-spp-sm2-1cls-scalexy_125_5000.weights --cfg ../yolov3-spp-sm2-1cls-scalexy_125.cfg --img-size 320 --conf-thres 0.2 # test 104 | python3 test.py --data ../supermarket2/supermarket2.data --weights weights/yolov3-spp-sm2-1cls-scalexy_150_5000.weights --cfg ../yolov3-spp-sm2-1cls-scalexy_150.cfg --img-size 320 --conf-thres 0.2 # test 105 | python3 test.py --data ../supermarket2/supermarket2.data --weights weights/yolov3-spp-sm2-1cls-scalexy_200_5000.weights --cfg ../yolov3-spp-sm2-1cls-scalexy_200.cfg --img-size 320 --conf-thres 0.2 # test 106 | python3 test.py --data ../supermarket2/supermarket2.data --weights ../darknet/backup/yolov3-spp-sm2-1cls-scalexy_variable_5000.weights --cfg ../yolov3-spp-sm2-1cls-scalexy_variable.cfg --img-size 320 --conf-thres 0.2 # test 107 | 108 | python3 train.py --img-size 320 --epochs 27 --batch-size 64 --accumulate 1 --nosave --notest && python3 test.py --weights weights/last.pt --img-size 320 --save-json && sudo shutdown 109 | 110 | # Debug/Development 111 | python3 train.py --data data/coco.data --img-size 320 --single-scale --batch-size 64 --accumulate 1 --epochs 1 --evolve --giou 112 | python3 test.py --weights weights/last.pt --cfg cfg/yolov3-spp.cfg --img-size 320 113 | 114 | gsutil cp evolve.txt gs://ultralytics 115 | sudo shutdown 116 | 117 | #Docker 118 | sudo docker kill $(sudo docker ps -q) 119 | sudo docker pull ultralytics/yolov3:v1 120 | sudo nvidia-docker run -it --ipc=host --mount type=bind,source="$(pwd)"/coco,target=/usr/src/coco ultralytics/yolov3:v1 121 | 122 | clear 123 | while true 124 | do 125 | python3 train.py --data data/coco.data --img-size 320 --batch-size 64 --accumulate 1 --evolve --epochs 1 --adam --bucket yolov4/adamdefaultpw_coco_1e --device 1 126 | done 127 | 128 | python3 train.py --data data/coco.data --img-size 320 --batch-size 64 --accumulate 1 --epochs 1 --adam --device 1 --prebias 129 | while true; do python3 train.py --data data/coco.data --img-size 320 --batch-size 64 --accumulate 1 --evolve --epochs 1 --adam --bucket yolov4/adamdefaultpw_coco_1e; done 130 | -------------------------------------------------------------------------------- /prune_yolov5s.py: -------------------------------------------------------------------------------- 1 | from modelsori import * 2 | from utils.utils import * 3 | import numpy as np 4 | from copy import deepcopy 5 | from test import test 6 | from terminaltables import AsciiTable 7 | import time 8 | from utils.prune_utils import * 9 | import argparse 10 | 11 | from utils.model_transfer import copy_weight_v6,copy_weight_v6x 12 | 13 | if __name__ == '__main__': 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('--cfg', type=str, default='cfg/yolov5s_v6_hand.cfg', help='cfg file path') 16 | parser.add_argument('--data', type=str, default='data/oxfordhand.data', help='*.data file path') 17 | parser.add_argument('--weights', type=str, default='weights/last_v6s_0.pt', help='sparse model weights') 18 | parser.add_argument('--percent', type=float, default=0.5, help='channel prune percent') 19 | parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') 20 | opt = parser.parse_args() 21 | print(opt) 22 | 23 | img_size = opt.img_size 24 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 25 | model = Darknet(opt.cfg, (img_size, img_size)).to(device) 26 | 27 | modelyolov5 = torch.load(opt.weights, map_location=device)['model'].float() # load FP32 model 28 | stride=32.0 29 | if len(modelyolov5.yaml["anchors"]) == 4: 30 | copy_weight_v6x(modelyolov5, model) 31 | stride=64.0 32 | else: 33 | copy_weight_v6(modelyolov5, model) 34 | 35 | eval_model = lambda model:test(opt.cfg, opt.data, 36 | weights=opt.weights, 37 | batch_size=4, 38 | img_size=img_size, 39 | iou_thres=0.5, 40 | conf_thres=0.001, 41 | nms_thres=0.5, 42 | stride=stride, 43 | save_json=False, 44 | model=model) 45 | obtain_num_parameters = lambda model:sum([param.nelement() for param in model.parameters()]) 46 | 47 | print("\nlet's test the original model first:") 48 | with torch.no_grad(): 49 | origin_model_metric = eval_model(model) 50 | 51 | origin_nparameters = obtain_num_parameters(model) 52 | 53 | CBL_idx, Conv_idx, prune_idx= parse_module_defs(model.module_defs) 54 | 55 | bn_weights = gather_bn_weights(model.module_list, prune_idx) 56 | 57 | sorted_bn = torch.sort(bn_weights)[0] 58 | 59 | # 避免剪掉所有channel的最高阈值(每个BN层的gamma的最大值的最小值即为阈值上限) 60 | highest_thre = [] 61 | for idx in prune_idx: 62 | # highest_thre.append(model.module_list[idx][1].weight.data.abs().max().item()) 63 | highest_thre.append(model.module_list[idx][1].weight.data.abs().max().item() if type(model.module_list[idx][1]).__name__ is 'BatchNorm2d' else model.module_list[idx][0].weight.data.abs().max().item()) 64 | highest_thre = min(highest_thre) 65 | 66 | # 找到highest_thre对应的下标对应的百分比 67 | if len((sorted_bn == highest_thre).nonzero()) > 1: 68 | high_num = (sorted_bn == highest_thre).nonzero()[1] 69 | else: 70 | high_num = (sorted_bn == highest_thre).nonzero() 71 | percent_limit = high_num.item() / len(bn_weights) 72 | # 优化好的模型只有一个极值点，上面代码临时使用以方便调试 73 | # 找到highest_thre对应的下标对应的百分比 74 | # percent_limit = (sorted_bn==highest_thre).nonzero().item()/len(bn_weights) 75 | 76 | print(f'Suggested Gamma threshold should be less than {highest_thre:.4f}.') 77 | print(f'The corresponding prune ratio is {percent_limit:.3f}, but you can set higher.') 78 | 79 | #%% 80 | def prune_and_eval(model, sorted_bn, percent=.0): 81 | model_copy = deepcopy(model) 82 | thre_index = int(len(sorted_bn) * percent) 83 | thre = sorted_bn[thre_index] 84 | 85 | print(f'Gamma value that less than {thre:.4f} are set to zero!') 86 | 87 | remain_num = 0 88 | for idx in prune_idx: 89 | bn_module = model_copy.module_list[idx][1] if type(model_copy.module_list[idx][1]).__name__ is 'BatchNorm2d' else model_copy.module_list[idx][0] 90 | 91 | mask = obtain_bn_mask(bn_module, thre) 92 | 93 | remain_num += int(mask.sum()) 94 | bn_module.weight.data.mul_(mask) 95 | print("let's test the current model!") 96 | with torch.no_grad(): 97 | mAP = eval_model(model_copy)[0][2] 98 | 99 | 100 | print(f'Number of channels has been reduced from {len(sorted_bn)} to {remain_num}') 101 | print(f'Prune ratio: {1-remain_num/len(sorted_bn):.3f}') 102 | print(f"mAP of the 'pruned' model is {mAP:.4f}") 103 | 104 | return thre 105 | 106 | percent = opt.percent 107 | print('the required prune percent is', percent) 108 | threshold = prune_and_eval(model, sorted_bn, percent) 109 | #%% 110 | def obtain_filters_mask(model, thre, CBL_idx, prune_idx): 111 | 112 | pruned = 0 113 | total = 0 114 | num_filters = [] 115 | filters_mask = [] 116 | for idx in CBL_idx: 117 | bn_module = model.module_list[idx][1] if type(model.module_list[idx][1]).__name__ is 'BatchNorm2d' else model.module_list[idx][0] 118 | if idx in prune_idx: 119 | 120 | mask = obtain_bn_mask(bn_module, thre).cpu().numpy() 121 | remain = int(mask.sum()) 122 | pruned = pruned + mask.shape[0] - remain 123 | 124 | if remain == 0: 125 | # print("Channels would be all pruned!") 126 | # raise Exception 127 | max_value = bn_module.weight.data.abs().max() 128 | mask = obtain_bn_mask(bn_module, max_value).cpu().numpy() 129 | remain = int(mask.sum()) 130 | pruned = pruned + mask.shape[0] - remain 131 | 132 | print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 133 | f'remaining channel: {remain:>4d}') 134 | else: 135 | mask = np.ones(bn_module.weight.data.shape) 136 | remain = mask.shape[0] 137 | 138 | total += mask.shape[0] 139 | num_filters.append(remain) 140 | filters_mask.append(mask.copy()) 141 | 142 | prune_ratio = pruned / total 143 | print(f'Prune channels: {pruned}\tPrune ratio: {prune_ratio:.3f}') 144 | 145 | return num_filters, filters_mask 146 | 147 | num_filters, filters_mask = obtain_filters_mask(model, threshold, CBL_idx, prune_idx) 148 | 149 | #%% 150 | CBLidx2mask = {idx: mask.astype('float32') for idx, mask in zip(CBL_idx, filters_mask)} 151 | 152 | pruned_model = prune_model_keep_size2(model, CBL_idx, CBL_idx, CBLidx2mask) 153 | 154 | print("\nnow prune the model but keep size,(actually add offset of BN beta to next layer), let's see how the mAP goes") 155 | with torch.no_grad(): 156 | eval_model(pruned_model) 157 | 158 | 159 | #%% 160 | compact_module_defs = deepcopy(model.module_defs) 161 | for idx, num in zip(CBL_idx, num_filters): 162 | assert compact_module_defs[idx]['type'] == 'convolutional' or compact_module_defs[idx]['type'] == 'convolutional_noconv' 163 | compact_module_defs[idx]['filters'] = str(num) 164 | 165 | #%% 166 | compact_model = Darknet([model.hyperparams.copy()] + compact_module_defs, (img_size, img_size)).to(device) 167 | compact_nparameters = obtain_num_parameters(compact_model) 168 | 169 | init_weights_from_loose_model(compact_model, pruned_model, CBL_idx, Conv_idx, CBLidx2mask) 170 | 171 | #%% 172 | random_input = torch.rand((1, 3, img_size, img_size)).to(device) 173 | 174 | def obtain_avg_forward_time(input, model, repeat=200): 175 | 176 | model.eval() 177 | start = time.time() 178 | with torch.no_grad(): 179 | for i in range(repeat): 180 | output = model(input)[0] 181 | avg_infer_time = (time.time() - start) / repeat 182 | 183 | return avg_infer_time, output 184 | 185 | print('\ntesting avg forward time...') 186 | pruned_forward_time, pruned_output = obtain_avg_forward_time(random_input, pruned_model) 187 | compact_forward_time, compact_output = obtain_avg_forward_time(random_input, compact_model) 188 | 189 | diff = (pruned_output-compact_output).abs().gt(0.001).sum().item() 190 | if diff > 0: 191 | print('Something wrong with the pruned model!') 192 | 193 | #%% 194 | # 在测试集上测试剪枝后的模型, 并统计模型的参数数量 195 | print('testing the mAP of final pruned model') 196 | with torch.no_grad(): 197 | compact_model_metric = eval_model(compact_model) 198 | 199 | 200 | #%% 201 | # 比较剪枝前后参数数量的变化、指标性能的变化 202 | metric_table = [ 203 | ["Metric", "Before", "After"], 204 | ["mAP", f'{origin_model_metric[0][2]:.6f}', f'{compact_model_metric[0][2]:.6f}'], 205 | ["Parameters", f"{origin_nparameters}", f"{compact_nparameters}"], 206 | ["Inference", f'{pruned_forward_time:.4f}', f'{compact_forward_time:.4f}'] 207 | ] 208 | print(AsciiTable(metric_table).table) 209 | 210 | #%% 211 | # 生成剪枝后的cfg文件并保存模型 212 | pruned_cfg_name = opt.cfg.replace('/', f'/prune_{percent}_') 213 | pruned_cfg_file = write_cfg(pruned_cfg_name, [model.hyperparams.copy()] + compact_module_defs) 214 | print(f'Config file has been saved: {pruned_cfg_file}') 215 | 216 | compact_model_name = opt.weights.replace('/', f'/prune_{percent}_') 217 | if compact_model_name.endswith('.pt'): 218 | chkpt = {'epoch': -1, 219 | 'best_fitness': None, 220 | 'training_results': None, 221 | 'model': compact_model.state_dict(), 222 | 'optimizer': None} 223 | torch.save(chkpt, compact_model_name) 224 | compact_model_name = compact_model_name.replace('.pt', '.weights') 225 | # save_weights(compact_model, compact_model_name) 226 | print(f'Compact model has been saved: {compact_model_name}') 227 | 228 | # def initialize_weights(model): 229 | # for m in model.modules(): 230 | # t = type(m) 231 | # if t is nn.Conv2d: 232 | # pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 233 | # elif t is nn.BatchNorm2d: 234 | # m.eps = 1e-3 235 | # m.momentum = 0.03 236 | # elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 237 | # m.inplace = True 238 | # 239 | # 240 | # model_load = Darknet('cfg/prune_0.8_yolov3-spp.cfg', (img_size, img_size)).to(device) 241 | # initialize_weights(model_load) 242 | # model_load.load_state_dict(torch.load('weights/converted.pt')['model']) 243 | # # load_darknet_weights(model_load, 'weights/prune_0.8_yolov3-spp-ultralytics.weights') 244 | # compact_forward_time, compact_output = obtain_avg_forward_time(random_input, compact_model) 245 | # load_forward_time, load_output = obtain_avg_forward_time(random_input, model_load) 246 | # 247 | # diff = (load_output - compact_output).abs().gt(0.001).sum().item() 248 | # if diff > 0: 249 | # print('Something wrong with the load model!') 250 | 251 | 252 | -------------------------------------------------------------------------------- /cfg/yolov5s_v6_hand.cfg: -------------------------------------------------------------------------------- 1 | [net] 2 | # Testing 3 | #batch=1 4 | #subdivisions=1 5 | # Training 6 | batch=64 7 | subdivisions=8 8 | width=608 9 | height=608 10 | channels=3 11 | momentum=0.949 12 | decay=0.0005 13 | angle=0 14 | saturation = 1.5 15 | exposure = 1.5 16 | hue=.1 17 | 18 | learning_rate=0.00261 19 | burn_in=1000 20 | max_batches = 500500 21 | policy=steps 22 | steps=400000,450000 23 | scales=.1,.1 24 | 25 | #cutmix=1 26 | mosaic=1 27 | 28 | #:104x104 54:52x52 85:26x26 104:13x13 for 416 29 | #yaml-0 30 | [convolutional] 31 | batch_normalize=1 32 | filters=32 33 | size=6 34 | stride=2 35 | pad=2 36 | activation=SiLU 37 | 38 | #yaml-1-Downsample 39 | [convolutional] 40 | batch_normalize=1 41 | filters=64 42 | size=3 43 | stride=2 44 | pad=1 45 | activation=SiLU 46 | 47 | #yaml-2-C3 48 | ## cv2 49 | [convolutional] 50 | batch_normalize=1 51 | filters=32 52 | size=1 53 | stride=1 54 | pad=1 55 | activation=SiLU 56 | 57 | [route] 58 | layers = -2 59 | 60 | ## cv1 61 | [convolutional] 62 | batch_normalize=1 63 | filters=32 64 | size=1 65 | stride=1 66 | pad=1 67 | activation=SiLU 68 | 69 | ## m-cv1 70 | [convolutional] 71 | batch_normalize=1 72 | filters=32 73 | size=1 74 | stride=1 75 | pad=1 76 | activation=SiLU 77 | 78 | ## m-cv2 79 | [convolutional] 80 | batch_normalize=1 81 | filters=32 82 | size=3 83 | stride=1 84 | pad=1 85 | activation=SiLU 86 | 87 | ## m-add 88 | [shortcut] 89 | from=-3 90 | activation=linear 91 | 92 | ## C3 cat 93 | [route] 94 | layers = -1,-6 95 | 96 | ## C3 cv3 97 | [convolutional] 98 | batch_normalize=1 99 | filters=64 100 | size=1 101 | stride=1 102 | pad=1 103 | activation=SiLU 104 | 105 | #yaml-3-Downsamples 106 | [convolutional] 107 | batch_normalize=1 108 | filters=128 109 | size=3 110 | stride=2 111 | pad=1 112 | activation=SiLU 113 | 114 | #yaml-4-C3 115 | ## cv2 116 | [convolutional] 117 | batch_normalize=1 118 | filters=64 119 | size=1 120 | stride=1 121 | pad=1 122 | activation=SiLU 123 | 124 | [route] 125 | layers = -2 126 | 127 | ## cv1 128 | [convolutional] 129 | batch_normalize=1 130 | filters=64 131 | size=1 132 | stride=1 133 | pad=1 134 | activation=SiLU 135 | 136 | ## m1-cv1 137 | [convolutional] 138 | batch_normalize=1 139 | filters=64 140 | size=1 141 | stride=1 142 | pad=1 143 | activation=SiLU 144 | 145 | ## m1-cv2 146 | [convolutional] 147 | batch_normalize=1 148 | filters=64 149 | size=3 150 | stride=1 151 | pad=1 152 | activation=SiLU 153 | 154 | ## m1-add 155 | [shortcut] 156 | from=-3 157 | activation=linear 158 | 159 | ## m2-cv1 160 | [convolutional] 161 | batch_normalize=1 162 | filters=64 163 | size=1 164 | stride=1 165 | pad=1 166 | activation=SiLU 167 | 168 | 169 | ## m2-cv2 170 | [convolutional] 171 | batch_normalize=1 172 | filters=64 173 | size=3 174 | stride=1 175 | pad=1 176 | activation=SiLU 177 | 178 | ## m2-add 179 | [shortcut] 180 | from=-3 181 | activation=linear 182 | 183 | ## C3-cat 184 | [route] 185 | layers = -1,-9 186 | 187 | ## C3-cv3 188 | [convolutional] 189 | batch_normalize=1 190 | filters=128 191 | size=1 192 | stride=1 193 | pad=1 194 | activation=SiLU 195 | 196 | # yaml-5-Downsample 197 | [convolutional] 198 | batch_normalize=1 199 | filters=256 200 | size=3 201 | stride=2 202 | pad=1 203 | activation=SiLU 204 | 205 | #yaml-6-C3 206 | [convolutional] 207 | batch_normalize=1 208 | filters=128 209 | size=1 210 | stride=1 211 | pad=1 212 | activation=SiLU 213 | 214 | [route] 215 | layers = -2 216 | 217 | [convolutional] 218 | batch_normalize=1 219 | filters=128 220 | size=1 221 | stride=1 222 | pad=1 223 | activation=SiLU 224 | 225 | [convolutional] 226 | batch_normalize=1 227 | filters=128 228 | size=1 229 | stride=1 230 | pad=1 231 | activation=SiLU 232 | 233 | [convolutional] 234 | batch_normalize=1 235 | filters=128 236 | size=3 237 | stride=1 238 | pad=1 239 | activation=SiLU 240 | 241 | [shortcut] 242 | from=-3 243 | activation=linear 244 | 245 | [convolutional] 246 | batch_normalize=1 247 | filters=128 248 | size=1 249 | stride=1 250 | pad=1 251 | activation=SiLU 252 | 253 | [convolutional] 254 | batch_normalize=1 255 | filters=128 256 | size=3 257 | stride=1 258 | pad=1 259 | activation=SiLU 260 | 261 | [shortcut] 262 | from=-3 263 | activation=linear 264 | 265 | [convolutional] 266 | batch_normalize=1 267 | filters=128 268 | size=1 269 | stride=1 270 | pad=1 271 | activation=SiLU 272 | 273 | [convolutional] 274 | batch_normalize=1 275 | filters=128 276 | size=3 277 | stride=1 278 | pad=1 279 | activation=SiLU 280 | 281 | [shortcut] 282 | from=-3 283 | activation=linear 284 | 285 | [route] 286 | layers = -1,-12 287 | 288 | [convolutional] 289 | batch_normalize=1 290 | filters=256 291 | size=1 292 | stride=1 293 | pad=1 294 | activation=SiLU 295 | 296 | # yaml-7-Downsample 297 | [convolutional] 298 | batch_normalize=1 299 | filters=512 300 | size=3 301 | stride=2 302 | pad=1 303 | activation=SiLU 304 | 305 | # yaml-8-C3 306 | ## cv2 307 | [convolutional] 308 | batch_normalize=1 309 | filters=256 310 | size=1 311 | stride=1 312 | pad=1 313 | activation=SiLU 314 | 315 | [route] 316 | layers = -2 317 | 318 | ## cv1 319 | [convolutional] 320 | batch_normalize=1 321 | filters=256 322 | size=1 323 | stride=1 324 | pad=1 325 | activation=SiLU 326 | 327 | ## m-cv1 328 | [convolutional] 329 | batch_normalize=1 330 | filters=256 331 | size=1 332 | stride=1 333 | pad=1 334 | activation=SiLU 335 | 336 | ## m-cv2 337 | [convolutional] 338 | batch_normalize=1 339 | filters=256 340 | size=3 341 | stride=1 342 | pad=1 343 | activation=SiLU 344 | 345 | ## m-add 346 | [shortcut] 347 | from=-3 348 | activation=linear 349 | 350 | ## C3 cat 351 | [route] 352 | layers = -1,-6 353 | 354 | ## C3 cv3 355 | [convolutional] 356 | batch_normalize=1 357 | filters=512 358 | size=1 359 | stride=1 360 | pad=1 361 | activation=SiLU 362 | 363 | #yaml-9-SPPF 364 | #SPPF cv1 365 | [convolutional] 366 | batch_normalize=1 367 | filters=256 368 | size=1 369 | stride=1 370 | pad=1 371 | activation=SiLU 372 | 373 | [maxpool] 374 | stride=1 375 | size=5 376 | 377 | [maxpool] 378 | stride=1 379 | size=5 380 | 381 | [maxpool] 382 | stride=1 383 | size=5 384 | 385 | [route] 386 | layers=-4,-3,-2,-1 387 | 388 | #SPPF cv2 389 | [convolutional] 390 | batch_normalize=1 391 | filters=512 392 | size=1 393 | stride=1 394 | pad=1 395 | activation=SiLU 396 | 397 | #yaml-10-conv 398 | [convolutional] 399 | batch_normalize=1 400 | filters=256 401 | size=1 402 | stride=1 403 | pad=1 404 | activation=SiLU 405 | 406 | #yaml-11-upsample 407 | [upsample] 408 | stride=2 409 | 410 | #yaml-12-concat 411 | [route] 412 | layers = -1,-18 413 | 414 | #yaml-13-C3 415 | [convolutional] 416 | batch_normalize=1 417 | filters=128 418 | size=1 419 | stride=1 420 | pad=1 421 | activation=SiLU 422 | 423 | [route] 424 | layers = -2 425 | 426 | [convolutional] 427 | batch_normalize=1 428 | filters=128 429 | size=1 430 | stride=1 431 | pad=1 432 | activation=SiLU 433 | 434 | [convolutional] 435 | batch_normalize=1 436 | filters=128 437 | size=1 438 | stride=1 439 | pad=1 440 | activation=SiLU 441 | 442 | [convolutional] 443 | batch_normalize=1 444 | filters=128 445 | size=3 446 | stride=1 447 | pad=1 448 | activation=SiLU 449 | 450 | [route] 451 | layers = -1,-5 452 | 453 | [convolutional] 454 | batch_normalize=1 455 | filters=256 456 | size=1 457 | stride=1 458 | pad=1 459 | activation=SiLU 460 | 461 | #yaml-14-conv 462 | [convolutional] 463 | batch_normalize=1 464 | filters=128 465 | size=1 466 | stride=1 467 | pad=1 468 | activation=SiLU 469 | 470 | #yaml-15-upsample 471 | [upsample] 472 | stride=2 473 | 474 | #yaml-16-concat 475 | [route] 476 | layers = -1,-43 477 | 478 | #yaml-17-C3 479 | [convolutional] 480 | batch_normalize=1 481 | filters=64 482 | size=1 483 | stride=1 484 | pad=1 485 | activation=SiLU 486 | 487 | [route] 488 | layers = -2 489 | 490 | [convolutional] 491 | batch_normalize=1 492 | filters=64 493 | size=1 494 | stride=1 495 | pad=1 496 | activation=SiLU 497 | 498 | [convolutional] 499 | batch_normalize=1 500 | filters=64 501 | size=1 502 | stride=1 503 | pad=1 504 | activation=SiLU 505 | 506 | [convolutional] 507 | batch_normalize=1 508 | filters=64 509 | size=3 510 | stride=1 511 | pad=1 512 | activation=SiLU 513 | 514 | [route] 515 | layers = -1,-5 516 | 517 | [convolutional] 518 | batch_normalize=1 519 | filters=128 520 | size=1 521 | stride=1 522 | pad=1 523 | activation=SiLU 524 | 525 | ##################### #filters=(classes + 5)x3 526 | [convolutional] 527 | size=1 528 | stride=1 529 | pad=1 530 | filters=18 531 | activation=linear 532 | 533 | [yolo] 534 | mask = 0,1,2 535 | anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326 536 | classes=1 537 | num=9 538 | jitter=.3 539 | ignore_thresh = .7 540 | truth_thresh = 1 541 | scale_x_y = 1.2 542 | iou_thresh=0.213 543 | cls_normalizer=1.0 544 | iou_normalizer=0.07 545 | iou_loss=ciou 546 | nms_kind=greedynms 547 | beta_nms=0.6 548 | 549 | [route] 550 | layers = -3 551 | 552 | #yaml-18-conv 553 | [convolutional] 554 | batch_normalize=1 555 | filters=128 556 | size=3 557 | stride=2 558 | pad=1 559 | activation=SiLU 560 | 561 | #yaml-19-concat 562 | [route] 563 | layers = -1,-14 564 | 565 | #yaml-20-C3 566 | [convolutional] 567 | batch_normalize=1 568 | filters=128 569 | size=1 570 | stride=1 571 | pad=1 572 | activation=SiLU 573 | 574 | [route] 575 | layers = -2 576 | 577 | [convolutional] 578 | batch_normalize=1 579 | filters=128 580 | size=1 581 | stride=1 582 | pad=1 583 | activation=SiLU 584 | 585 | [convolutional] 586 | batch_normalize=1 587 | filters=128 588 | size=1 589 | stride=1 590 | pad=1 591 | activation=SiLU 592 | 593 | [convolutional] 594 | batch_normalize=1 595 | filters=128 596 | size=3 597 | stride=1 598 | pad=1 599 | activation=SiLU 600 | 601 | [route] 602 | layers = -1,-5 603 | 604 | [convolutional] 605 | batch_normalize=1 606 | filters=256 607 | size=1 608 | stride=1 609 | pad=1 610 | activation=SiLU 611 | 612 | ##################### #filters=(classes + 5)x3 613 | [convolutional] 614 | size=1 615 | stride=1 616 | pad=1 617 | filters=18 618 | activation=linear 619 | 620 | [yolo] 621 | mask = 3,4,5 622 | anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326 623 | classes=1 624 | num=9 625 | jitter=.3 626 | ignore_thresh = .7 627 | truth_thresh = 1 628 | scale_x_y = 1.2 629 | iou_thresh=0.213 630 | cls_normalizer=1.0 631 | iou_normalizer=0.07 632 | iou_loss=ciou 633 | nms_kind=greedynms 634 | beta_nms=0.6 635 | 636 | [route] 637 | layers = -3 638 | 639 | #yaml-21-conv 640 | [convolutional] 641 | batch_normalize=1 642 | filters=256 643 | size=3 644 | stride=2 645 | pad=1 646 | activation=SiLU 647 | 648 | #yaml-22-concat 649 | [route] 650 | layers = -1,-36 651 | 652 | #yaml-23-C3 653 | [convolutional] 654 | batch_normalize=1 655 | filters=256 656 | size=1 657 | stride=1 658 | pad=1 659 | activation=SiLU 660 | 661 | [route] 662 | layers = -2 663 | 664 | [convolutional] 665 | batch_normalize=1 666 | filters=256 667 | size=1 668 | stride=1 669 | pad=1 670 | activation=SiLU 671 | 672 | [convolutional] 673 | batch_normalize=1 674 | filters=256 675 | size=1 676 | stride=1 677 | pad=1 678 | activation=SiLU 679 | 680 | [convolutional] 681 | batch_normalize=1 682 | filters=256 683 | size=3 684 | stride=1 685 | pad=1 686 | activation=SiLU 687 | 688 | [route] 689 | layers = -1,-5 690 | 691 | [convolutional] 692 | batch_normalize=1 693 | filters=512 694 | size=1 695 | stride=1 696 | pad=1 697 | activation=SiLU 698 | 699 | ##################### #filters=(classes + 5)x3 700 | [convolutional] 701 | size=1 702 | stride=1 703 | pad=1 704 | filters=18 705 | activation=linear 706 | 707 | [yolo] 708 | mask = 6,7,8 709 | anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326 710 | classes=1 711 | num=9 712 | jitter=.3 713 | ignore_thresh = .7 714 | truth_thresh = 1 715 | scale_x_y = 1.2 716 | iou_thresh=0.213 717 | cls_normalizer=1.0 718 | iou_normalizer=0.07 719 | iou_loss=ciou 720 | nms_kind=greedynms 721 | beta_nms=0.6 722 | 723 | 724 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /utils/adabound.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | from torch.optim import Optimizer 5 | 6 | 7 | class AdaBound(Optimizer): 8 | """Implements AdaBound algorithm. 9 | It has been proposed in `Adaptive Gradient Methods with Dynamic Bound of Learning Rate`_. 10 | Arguments: 11 | params (iterable): iterable of parameters to optimize or dicts defining 12 | parameter groups 13 | lr (float, optional): Adam learning rate (default: 1e-3) 14 | betas (Tuple[float, float], optional): coefficients used for computing 15 | running averages of gradient and its square (default: (0.9, 0.999)) 16 | final_lr (float, optional): final (SGD) learning rate (default: 0.1) 17 | gamma (float, optional): convergence speed of the bound functions (default: 1e-3) 18 | eps (float, optional): term added to the denominator to improve 19 | numerical stability (default: 1e-8) 20 | weight_decay (float, optional): weight decay (L2 penalty) (default: 0) 21 | amsbound (boolean, optional): whether to use the AMSBound variant of this algorithm 22 | .. Adaptive Gradient Methods with Dynamic Bound of Learning Rate: 23 | https://openreview.net/forum?id=Bkg3g2R9FX 24 | """ 25 | 26 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), final_lr=0.1, gamma=1e-3, 27 | eps=1e-8, weight_decay=0, amsbound=False): 28 | if not 0.0 <= lr: 29 | raise ValueError("Invalid learning rate: {}".format(lr)) 30 | if not 0.0 <= eps: 31 | raise ValueError("Invalid epsilon value: {}".format(eps)) 32 | if not 0.0 <= betas[0] < 1.0: 33 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 34 | if not 0.0 <= betas[1] < 1.0: 35 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 36 | if not 0.0 <= final_lr: 37 | raise ValueError("Invalid final learning rate: {}".format(final_lr)) 38 | if not 0.0 <= gamma < 1.0: 39 | raise ValueError("Invalid gamma parameter: {}".format(gamma)) 40 | defaults = dict(lr=lr, betas=betas, final_lr=final_lr, gamma=gamma, eps=eps, 41 | weight_decay=weight_decay, amsbound=amsbound) 42 | super(AdaBound, self).__init__(params, defaults) 43 | 44 | self.base_lrs = list(map(lambda group: group['lr'], self.param_groups)) 45 | 46 | def __setstate__(self, state): 47 | super(AdaBound, self).__setstate__(state) 48 | for group in self.param_groups: 49 | group.setdefault('amsbound', False) 50 | 51 | def step(self, closure=None): 52 | """Performs a single optimization step. 53 | Arguments: 54 | closure (callable, optional): A closure that reevaluates the model 55 | and returns the loss. 56 | """ 57 | loss = None 58 | if closure is not None: 59 | loss = closure() 60 | 61 | for group, base_lr in zip(self.param_groups, self.base_lrs): 62 | for p in group['params']: 63 | if p.grad is None: 64 | continue 65 | grad = p.grad.data 66 | if grad.is_sparse: 67 | raise RuntimeError( 68 | 'Adam does not support sparse gradients, please consider SparseAdam instead') 69 | amsbound = group['amsbound'] 70 | 71 | state = self.state[p] 72 | 73 | # State initialization 74 | if len(state) == 0: 75 | state['step'] = 0 76 | # Exponential moving average of gradient values 77 | state['exp_avg'] = torch.zeros_like(p.data) 78 | # Exponential moving average of squared gradient values 79 | state['exp_avg_sq'] = torch.zeros_like(p.data) 80 | if amsbound: 81 | # Maintains max of all exp. moving avg. of sq. grad. values 82 | state['max_exp_avg_sq'] = torch.zeros_like(p.data) 83 | 84 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 85 | if amsbound: 86 | max_exp_avg_sq = state['max_exp_avg_sq'] 87 | beta1, beta2 = group['betas'] 88 | 89 | state['step'] += 1 90 | 91 | if group['weight_decay'] != 0: 92 | grad = grad.add(group['weight_decay'], p.data) 93 | 94 | # Decay the first and second moment running average coefficient 95 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 96 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 97 | if amsbound: 98 | # Maintains the maximum of all 2nd moment running avg. till now 99 | torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) 100 | # Use the max. for normalizing running avg. of gradient 101 | denom = max_exp_avg_sq.sqrt().add_(group['eps']) 102 | else: 103 | denom = exp_avg_sq.sqrt().add_(group['eps']) 104 | 105 | bias_correction1 = 1 - beta1 ** state['step'] 106 | bias_correction2 = 1 - beta2 ** state['step'] 107 | step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 108 | 109 | # Applies bounds on actual learning rate 110 | # lr_scheduler cannot affect final_lr, this is a workaround to apply lr decay 111 | final_lr = group['final_lr'] * group['lr'] / base_lr 112 | lower_bound = final_lr * (1 - 1 / (group['gamma'] * state['step'] + 1)) 113 | upper_bound = final_lr * (1 + 1 / (group['gamma'] * state['step'])) 114 | step_size = torch.full_like(denom, step_size) 115 | step_size.div_(denom).clamp_(lower_bound, upper_bound).mul_(exp_avg) 116 | 117 | p.data.add_(-step_size) 118 | 119 | return loss 120 | 121 | 122 | class AdaBoundW(Optimizer): 123 | """Implements AdaBound algorithm with Decoupled Weight Decay (arxiv.org/abs/1711.05101) 124 | It has been proposed in `Adaptive Gradient Methods with Dynamic Bound of Learning Rate`_. 125 | Arguments: 126 | params (iterable): iterable of parameters to optimize or dicts defining 127 | parameter groups 128 | lr (float, optional): Adam learning rate (default: 1e-3) 129 | betas (Tuple[float, float], optional): coefficients used for computing 130 | running averages of gradient and its square (default: (0.9, 0.999)) 131 | final_lr (float, optional): final (SGD) learning rate (default: 0.1) 132 | gamma (float, optional): convergence speed of the bound functions (default: 1e-3) 133 | eps (float, optional): term added to the denominator to improve 134 | numerical stability (default: 1e-8) 135 | weight_decay (float, optional): weight decay (L2 penalty) (default: 0) 136 | amsbound (boolean, optional): whether to use the AMSBound variant of this algorithm 137 | .. Adaptive Gradient Methods with Dynamic Bound of Learning Rate: 138 | https://openreview.net/forum?id=Bkg3g2R9FX 139 | """ 140 | 141 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), final_lr=0.1, gamma=1e-3, 142 | eps=1e-8, weight_decay=0, amsbound=False): 143 | if not 0.0 <= lr: 144 | raise ValueError("Invalid learning rate: {}".format(lr)) 145 | if not 0.0 <= eps: 146 | raise ValueError("Invalid epsilon value: {}".format(eps)) 147 | if not 0.0 <= betas[0] < 1.0: 148 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 149 | if not 0.0 <= betas[1] < 1.0: 150 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 151 | if not 0.0 <= final_lr: 152 | raise ValueError("Invalid final learning rate: {}".format(final_lr)) 153 | if not 0.0 <= gamma < 1.0: 154 | raise ValueError("Invalid gamma parameter: {}".format(gamma)) 155 | defaults = dict(lr=lr, betas=betas, final_lr=final_lr, gamma=gamma, eps=eps, 156 | weight_decay=weight_decay, amsbound=amsbound) 157 | super(AdaBoundW, self).__init__(params, defaults) 158 | 159 | self.base_lrs = list(map(lambda group: group['lr'], self.param_groups)) 160 | 161 | def __setstate__(self, state): 162 | super(AdaBoundW, self).__setstate__(state) 163 | for group in self.param_groups: 164 | group.setdefault('amsbound', False) 165 | 166 | def step(self, closure=None): 167 | """Performs a single optimization step. 168 | Arguments: 169 | closure (callable, optional): A closure that reevaluates the model 170 | and returns the loss. 171 | """ 172 | loss = None 173 | if closure is not None: 174 | loss = closure() 175 | 176 | for group, base_lr in zip(self.param_groups, self.base_lrs): 177 | for p in group['params']: 178 | if p.grad is None: 179 | continue 180 | grad = p.grad.data 181 | if grad.is_sparse: 182 | raise RuntimeError( 183 | 'Adam does not support sparse gradients, please consider SparseAdam instead') 184 | amsbound = group['amsbound'] 185 | 186 | state = self.state[p] 187 | 188 | # State initialization 189 | if len(state) == 0: 190 | state['step'] = 0 191 | # Exponential moving average of gradient values 192 | state['exp_avg'] = torch.zeros_like(p.data) 193 | # Exponential moving average of squared gradient values 194 | state['exp_avg_sq'] = torch.zeros_like(p.data) 195 | if amsbound: 196 | # Maintains max of all exp. moving avg. of sq. grad. values 197 | state['max_exp_avg_sq'] = torch.zeros_like(p.data) 198 | 199 | exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] 200 | if amsbound: 201 | max_exp_avg_sq = state['max_exp_avg_sq'] 202 | beta1, beta2 = group['betas'] 203 | 204 | state['step'] += 1 205 | 206 | # Decay the first and second moment running average coefficient 207 | exp_avg.mul_(beta1).add_(1 - beta1, grad) 208 | exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) 209 | if amsbound: 210 | # Maintains the maximum of all 2nd moment running avg. till now 211 | torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) 212 | # Use the max. for normalizing running avg. of gradient 213 | denom = max_exp_avg_sq.sqrt().add_(group['eps']) 214 | else: 215 | denom = exp_avg_sq.sqrt().add_(group['eps']) 216 | 217 | bias_correction1 = 1 - beta1 ** state['step'] 218 | bias_correction2 = 1 - beta2 ** state['step'] 219 | step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 220 | 221 | # Applies bounds on actual learning rate 222 | # lr_scheduler cannot affect final_lr, this is a workaround to apply lr decay 223 | final_lr = group['final_lr'] * group['lr'] / base_lr 224 | lower_bound = final_lr * (1 - 1 / (group['gamma'] * state['step'] + 1)) 225 | upper_bound = final_lr * (1 + 1 / (group['gamma'] * state['step'])) 226 | step_size = torch.full_like(denom, step_size) 227 | step_size.div_(denom).clamp_(lower_bound, upper_bound).mul_(exp_avg) 228 | 229 | if group['weight_decay'] != 0: 230 | decayed_weights = torch.mul(p.data, group['weight_decay']) 231 | p.data.add_(-step_size) 232 | p.data.sub_(decayed_weights) 233 | else: 234 | p.data.add_(-step_size) 235 | 236 | return loss 237 | -------------------------------------------------------------------------------- /data/coco_128img.txt: -------------------------------------------------------------------------------- 1 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000009.jpg 2 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000025.jpg 3 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000030.jpg 4 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000034.jpg 5 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000036.jpg 6 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000042.jpg 7 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000049.jpg 8 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000061.jpg 9 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000064.jpg 10 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000071.jpg 11 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000072.jpg 12 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000073.jpg 13 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000074.jpg 14 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000077.jpg 15 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000078.jpg 16 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000081.jpg 17 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000086.jpg 18 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000089.jpg 19 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000092.jpg 20 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000094.jpg 21 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000109.jpg 22 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000110.jpg 23 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000113.jpg 24 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000127.jpg 25 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000133.jpg 26 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000136.jpg 27 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000138.jpg 28 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000142.jpg 29 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000143.jpg 30 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000144.jpg 31 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000149.jpg 32 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000151.jpg 33 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000154.jpg 34 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000164.jpg 35 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000165.jpg 36 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000192.jpg 37 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000194.jpg 38 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000196.jpg 39 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000201.jpg 40 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000208.jpg 41 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000241.jpg 42 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000247.jpg 43 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000250.jpg 44 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000257.jpg 45 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000260.jpg 46 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000263.jpg 47 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000283.jpg 48 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000294.jpg 49 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000307.jpg 50 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000308.jpg 51 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000309.jpg 52 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000312.jpg 53 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000315.jpg 54 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000321.jpg 55 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000322.jpg 56 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000326.jpg 57 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000328.jpg 58 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000332.jpg 59 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000338.jpg 60 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000349.jpg 61 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000357.jpg 62 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000359.jpg 63 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000360.jpg 64 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000368.jpg 65 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000370.jpg 66 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000382.jpg 67 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000384.jpg 68 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000387.jpg 69 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000389.jpg 70 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000394.jpg 71 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000395.jpg 72 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000397.jpg 73 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000400.jpg 74 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000404.jpg 75 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000415.jpg 76 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000419.jpg 77 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000428.jpg 78 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000431.jpg 79 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000436.jpg 80 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000438.jpg 81 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000443.jpg 82 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000446.jpg 83 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000450.jpg 84 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000459.jpg 85 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000471.jpg 86 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000472.jpg 87 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000474.jpg 88 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000486.jpg 89 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000488.jpg 90 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000490.jpg 91 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000491.jpg 92 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000502.jpg 93 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000508.jpg 94 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000510.jpg 95 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000514.jpg 96 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000520.jpg 97 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000529.jpg 98 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000531.jpg 99 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000532.jpg 100 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000536.jpg 101 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000540.jpg 102 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000542.jpg 103 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000544.jpg 104 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000560.jpg 105 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000562.jpg 106 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000564.jpg 107 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000569.jpg 108 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000572.jpg 109 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000575.jpg 110 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000581.jpg 111 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000584.jpg 112 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000589.jpg 113 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000590.jpg 114 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000595.jpg 115 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000597.jpg 116 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000599.jpg 117 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000605.jpg 118 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000612.jpg 119 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000620.jpg 120 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000623.jpg 121 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000625.jpg 122 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000626.jpg 123 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000629.jpg 124 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000634.jpg 125 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000636.jpg 126 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000641.jpg 127 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000643.jpg 128 | /home/lishuang/Disk/remote/pycharm/yolov5_prune/data/coco128/images/train2017/000000000650.jpg -------------------------------------------------------------------------------- /slim_prune_yolov5s.py: -------------------------------------------------------------------------------- 1 | from modelsori import * 2 | from utils.utils import * 3 | import numpy as np 4 | from copy import deepcopy 5 | from test import test 6 | from terminaltables import AsciiTable 7 | import time 8 | from utils.prune_utils import * 9 | import argparse 10 | import torchvision 11 | 12 | from utils.model_transfer import copy_weight_v6,copy_weight_v6x 13 | 14 | if __name__ == '__main__': 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--cfg', type=str, default='cfg/yolov5s_v6_hand.cfg', help='cfg file path') 17 | parser.add_argument('--data', type=str, default='data/oxfordhand.data', help='*.data file path') 18 | parser.add_argument('--weights', type=str, default='weights/last_v6s.pt', help='sparse model weights') 19 | parser.add_argument('--global_percent', type=float, default=0.6, help='global channel prune percent') 20 | parser.add_argument('--layer_keep', type=float, default=0.01, help='channel keep percent per layer') 21 | parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') 22 | opt = parser.parse_args() 23 | print(opt) 24 | 25 | img_size = opt.img_size 26 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 27 | model = Darknet(opt.cfg, (img_size, img_size)).to(device) 28 | 29 | modelyolov5 = torch.load(opt.weights, map_location=device)['model'].float() # load FP32 model 30 | stride=32.0 31 | if len(modelyolov5.yaml["anchors"]) == 4: 32 | copy_weight_v6x(modelyolov5, model) 33 | stride=64.0 34 | else: 35 | copy_weight_v6(modelyolov5, model) 36 | 37 | # model.load_state_dict(torch.load(opt.weights)['model'].state_dict()) 38 | 39 | 40 | eval_model = lambda model:test(model=model,cfg=opt.cfg, data=opt.data, batch_size=4, img_size=img_size,stride=stride) 41 | obtain_num_parameters = lambda model:sum([param.nelement() for param in model.parameters()]) 42 | 43 | print("\nlet's test the original model first:") 44 | with torch.no_grad(): 45 | origin_model_metric = eval_model(model) 46 | origin_nparameters = obtain_num_parameters(model) 47 | 48 | CBL_idx, Conv_idx, prune_idx, _, _= parse_module_defs2(model.module_defs) 49 | 50 | 51 | 52 | bn_weights = gather_bn_weights(model.module_list, prune_idx) 53 | 54 | sorted_bn = torch.sort(bn_weights)[0] 55 | sorted_bn, sorted_index = torch.sort(bn_weights) 56 | thresh_index = int(len(bn_weights) * opt.global_percent) 57 | thresh = sorted_bn[thresh_index].cuda() 58 | 59 | print(f'Global Threshold should be less than {thresh:.4f}.') 60 | 61 | 62 | 63 | 64 | #%% 65 | def obtain_filters_mask(model, thre, CBL_idx, prune_idx): 66 | 67 | pruned = 0 68 | total = 0 69 | num_filters = [] 70 | filters_mask = [] 71 | for idx in CBL_idx: 72 | # bn_module = model.module_list[idx][1] 73 | bn_module = model.module_list[idx][1] \ 74 | if type(model.module_list[idx][1]).__name__ == 'BatchNorm2d' \ 75 | else model.module_list[idx][0] 76 | if idx in prune_idx: 77 | 78 | weight_copy = bn_module.weight.data.abs().clone() 79 | 80 | if model.module_defs[idx]['type'] == 'convolutional_noconv': 81 | channels = weight_copy.shape[0] 82 | channels_half = int(channels / 2) 83 | weight_copy1 = weight_copy[:channels_half] 84 | weight_copy2 = weight_copy[channels_half:] 85 | min_channel_num = int(channels_half * opt.layer_keep) if int(channels_half * opt.layer_keep) > 0 else 1 86 | mask1 = weight_copy1.gt(thresh).float() 87 | mask2 = weight_copy2.gt(thresh).float() 88 | 89 | if int(torch.sum(mask1)) < min_channel_num: 90 | _, sorted_index_weights1 = torch.sort(weight_copy1, descending=True) 91 | mask1[sorted_index_weights1[:min_channel_num]] = 1. 92 | 93 | if int(torch.sum(mask2)) < min_channel_num: 94 | _, sorted_index_weights2 = torch.sort(weight_copy2, descending=True) 95 | mask2[sorted_index_weights2[:min_channel_num]] = 1. 96 | 97 | remain1 = int(mask1.sum()) 98 | pruned = pruned + mask1.shape[0] - remain1 99 | remain2 = int(mask2.sum()) 100 | pruned = pruned + mask2.shape[0] - remain2 101 | 102 | mask = torch.cat((mask1, mask2)) 103 | remain = remain1 + remain2 104 | 105 | print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 106 | f'remaining channel: {remain:>4d}') 107 | else: 108 | 109 | channels = weight_copy.shape[0] # 110 | min_channel_num = int(channels * opt.layer_keep) if int(channels * opt.layer_keep) > 0 else 1 111 | mask = weight_copy.gt(thresh).float() 112 | 113 | if int(torch.sum(mask)) < min_channel_num: 114 | _, sorted_index_weights = torch.sort(weight_copy, descending=True) 115 | mask[sorted_index_weights[:min_channel_num]] = 1. 116 | 117 | remain = int(mask.sum()) 118 | pruned = pruned + mask.shape[0] - remain 119 | 120 | print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 121 | f'remaining channel: {remain:>4d}') 122 | ########### 123 | 124 | # channels = weight_copy.shape[0] # 125 | # min_channel_num = int(channels * opt.layer_keep) if int(channels * opt.layer_keep) > 0 else 1 126 | # mask = weight_copy.gt(thresh).float() 127 | # 128 | # if int(torch.sum(mask)) < min_channel_num: 129 | # _, sorted_index_weights = torch.sort(weight_copy,descending=True) 130 | # mask[sorted_index_weights[:min_channel_num]]=1. 131 | # remain = int(mask.sum()) 132 | # pruned = pruned + mask.shape[0] - remain 133 | # 134 | # print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 135 | # f'remaining channel: {remain:>4d}') 136 | else: 137 | mask = torch.ones(bn_module.weight.data.shape) 138 | remain = mask.shape[0] 139 | 140 | total += mask.shape[0] 141 | num_filters.append(remain) 142 | filters_mask.append(mask.clone()) 143 | 144 | prune_ratio = pruned / total 145 | print(f'Prune channels: {pruned}\tPrune ratio: {prune_ratio:.3f}') 146 | 147 | return num_filters, filters_mask 148 | 149 | num_filters, filters_mask = obtain_filters_mask(model, thresh, CBL_idx, prune_idx) 150 | CBLidx2mask = {idx: mask for idx, mask in zip(CBL_idx, filters_mask)} 151 | CBLidx2filters = {idx: filters for idx, filters in zip(CBL_idx, num_filters)} 152 | 153 | for i in model.module_defs: 154 | if i['type'] == 'shortcut': 155 | i['is_access'] = False 156 | 157 | print('merge the mask of layers connected to shortcut!') 158 | merge_mask(model, CBLidx2mask, CBLidx2filters) 159 | 160 | 161 | 162 | def prune_and_eval(model, CBL_idx, CBLidx2mask): 163 | model_copy = deepcopy(model) 164 | 165 | for idx in CBL_idx: 166 | # bn_module = model_copy.module_list[idx][1] 167 | bn_module = model_copy.module_list[idx][1] \ 168 | if type(model_copy.module_list[idx][1]).__name__ == 'BatchNorm2d' \ 169 | else model_copy.module_list[idx][0] 170 | mask = CBLidx2mask[idx].cuda() 171 | bn_module.weight.data.mul_(mask) 172 | 173 | with torch.no_grad(): 174 | mAP = eval_model(model_copy)[0][2] 175 | 176 | print(f'mask the gamma as zero, mAP of the model is {mAP:.4f}') 177 | 178 | 179 | prune_and_eval(model, CBL_idx, CBLidx2mask) 180 | 181 | 182 | for i in CBLidx2mask: 183 | CBLidx2mask[i] = CBLidx2mask[i].clone().cpu().numpy() 184 | 185 | 186 | 187 | pruned_model = prune_model_keep_size2(model, prune_idx, CBL_idx, CBLidx2mask) 188 | print("\nnow prune the model but keep size,(actually add offset of BN beta to following layers), let's see how the mAP goes") 189 | 190 | with torch.no_grad(): 191 | eval_model(pruned_model) 192 | 193 | for i in model.module_defs: 194 | if i['type'] == 'shortcut': 195 | i.pop('is_access') 196 | 197 | # compact_module_defs = deepcopy(model.module_defs) 198 | # for idx in CBL_idx: 199 | # assert compact_module_defs[idx]['type'] == 'convolutional' 200 | # compact_module_defs[idx]['filters'] = str(CBLidx2filters[idx]) 201 | 202 | compact_module_defs = deepcopy(model.module_defs) 203 | for idx in CBL_idx: 204 | assert compact_module_defs[idx]['type'] == 'convolutional' or compact_module_defs[idx][ 205 | 'type'] == 'convolutional_noconv' 206 | num=CBLidx2filters[idx] 207 | compact_module_defs[idx]['filters'] = str(num) 208 | if compact_module_defs[idx]['type'] == 'convolutional_noconv': 209 | model_def = compact_module_defs[idx - 1] # route 210 | assert compact_module_defs[idx - 1]['type'] == 'route' 211 | from_layers = [int(s) for s in model_def['layers'].split(',')] 212 | assert compact_module_defs[idx - 1 + from_layers[0]]['type'] == 'convolutional_nobias' 213 | assert compact_module_defs[idx - 1 + from_layers[1] if from_layers[1] < 0 else from_layers[1]][ 214 | 'type'] == 'convolutional_nobias' 215 | half_num = int(len(CBLidx2mask[idx]) / 2) 216 | mask1 = CBLidx2mask[idx][:half_num] 217 | mask2 = CBLidx2mask[idx][half_num:] 218 | remain1 = int(mask1.sum()) 219 | remain2 = int(mask2.sum()) 220 | compact_module_defs[idx - 1 + from_layers[0]]['filters'] = remain1 221 | compact_module_defs[idx - 1 + from_layers[1] if from_layers[1] < 0 else from_layers[1]]['filters'] = remain2 222 | 223 | 224 | compact_model = Darknet([model.hyperparams.copy()] + compact_module_defs, (img_size, img_size)).to(device) 225 | compact_nparameters = obtain_num_parameters(compact_model) 226 | 227 | init_weights_from_loose_model(compact_model, pruned_model, CBL_idx, Conv_idx, CBLidx2mask) 228 | 229 | 230 | random_input = torch.rand((1, 3, img_size, img_size)).to(device) 231 | 232 | def obtain_avg_forward_time(input, model, repeat=200): 233 | # model.to('cpu').fuse() 234 | # model.module_list.to(device) 235 | model.eval() 236 | start = time.time() 237 | with torch.no_grad(): 238 | for i in range(repeat): 239 | output = model(input)[0] 240 | avg_infer_time = (time.time() - start) / repeat 241 | 242 | return avg_infer_time, output 243 | 244 | print('testing inference time...') 245 | pruned_forward_time, pruned_output = obtain_avg_forward_time(random_input, pruned_model) 246 | compact_forward_time, compact_output = obtain_avg_forward_time(random_input, compact_model) 247 | 248 | diff = (pruned_output - compact_output).abs().gt(0.001).sum().item() 249 | if diff > 0: 250 | print('Something wrong with the pruned model!') 251 | 252 | print('testing the final model...') 253 | with torch.no_grad(): 254 | compact_model_metric = eval_model(compact_model) 255 | 256 | 257 | metric_table = [ 258 | ["Metric", "Before", "After"], 259 | ["mAP", f'{origin_model_metric[0][2]:.6f}', f'{compact_model_metric[0][2]:.6f}'], 260 | ["Parameters", f"{origin_nparameters}", f"{compact_nparameters}"], 261 | ["Inference", f'{pruned_forward_time:.4f}', f'{compact_forward_time:.4f}'] 262 | ] 263 | print(AsciiTable(metric_table).table) 264 | 265 | 266 | 267 | pruned_cfg_name = opt.cfg.replace('/', f'/prune_{opt.global_percent}_keep_{opt.layer_keep}_') 268 | pruned_cfg_file = write_cfg(pruned_cfg_name, [model.hyperparams.copy()] + compact_module_defs) 269 | print(f'Config file has been saved: {pruned_cfg_file}') 270 | 271 | compact_model_name = opt.weights.replace('/', f'/prune_{opt.global_percent}_keep_{opt.layer_keep}_') 272 | if compact_model_name.endswith('.pt'): 273 | chkpt = {'epoch': -1, 274 | 'best_fitness': None, 275 | 'training_results': None, 276 | 'model': compact_model.state_dict(), 277 | 'optimizer': None} 278 | torch.save(chkpt, compact_model_name) 279 | compact_model_name = compact_model_name.replace('.pt', '.weights') 280 | # save_weights(compact_model, path=compact_model_name) 281 | print(f'Compact model has been saved: {compact_model_name}') 282 | 283 | -------------------------------------------------------------------------------- /shortcut_prune_yolov5s.py: -------------------------------------------------------------------------------- 1 | from modelsori import * 2 | from utils.utils import * 3 | import numpy as np 4 | from copy import deepcopy 5 | from test import test 6 | from terminaltables import AsciiTable 7 | import time 8 | from utils.prune_utils import * 9 | import argparse 10 | 11 | from utils.model_transfer import copy_weight_v6,copy_weight_v6x 12 | 13 | if __name__ == '__main__': 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('--cfg', type=str, default='cfg/yolov5s_v6_hand.cfg', help='cfg file path') 16 | parser.add_argument('--data', type=str, default='data/oxfordhand.data', help='*.data file path') 17 | parser.add_argument('--weights', type=str, default='weights/last_v6s.pt', help='sparse model weights') 18 | parser.add_argument('--percent', type=float, default=0.6, help='channel prune percent') 19 | parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') 20 | opt = parser.parse_args() 21 | print(opt) 22 | 23 | img_size = opt.img_size 24 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 25 | model = Darknet(opt.cfg, (img_size, img_size)).to(device) 26 | 27 | modelyolov5 = torch.load(opt.weights, map_location=device)['model'].float() # load FP32 model 28 | stride=32.0 29 | if len(modelyolov5.yaml["anchors"]) == 4: 30 | copy_weight_v6x(modelyolov5, model) 31 | stride=64.0 32 | else: 33 | copy_weight_v6(modelyolov5, model) 34 | 35 | eval_model = lambda model:test(model=model,cfg=opt.cfg, data=opt.data, batch_size=4, img_size=img_size,stride=stride) 36 | obtain_num_parameters = lambda model:sum([param.nelement() for param in model.parameters()]) 37 | 38 | print("\nlet's test the original model first:") 39 | with torch.no_grad(): 40 | origin_model_metric = eval_model(model) 41 | origin_nparameters = obtain_num_parameters(model) 42 | 43 | CBL_idx, Conv_idx, prune_idx,shortcut_idx,shortcut_all= parse_module_defs2(model.module_defs) 44 | 45 | 46 | sort_prune_idx=[idx for idx in prune_idx if idx not in shortcut_idx] 47 | 48 | #将所有要剪枝的BN层的α参数，拷贝到bn_weights列表 49 | bn_weights = gather_bn_weights(model.module_list, sort_prune_idx) 50 | 51 | #torch.sort返回二维列表，第一维是排序后的值列表，第二维是排序后的值列表对应的索引 52 | sorted_bn = torch.sort(bn_weights)[0] 53 | 54 | 55 | #避免剪掉所有channel的最高阈值(每个BN层的gamma的最大值的最小值即为阈值上限) 56 | highest_thre = [] 57 | for idx in sort_prune_idx: 58 | #.item()可以得到张量里的元素值 59 | # highest_thre.append(model.module_list[idx][1].weight.data.abs().max().item()) 60 | highest_thre.append(model.module_list[idx][1].weight.data.abs().max().item() if type( 61 | model.module_list[idx][1]).__name__ is 'BatchNorm2d' else model.module_list[idx][ 62 | 0].weight.data.abs().max().item()) 63 | highest_thre = min(highest_thre) 64 | 65 | # 找到highest_thre对应的下标对应的百分比 66 | if len((sorted_bn==highest_thre).nonzero()) > 1: 67 | high_num=(sorted_bn==highest_thre).nonzero()[1] 68 | else: 69 | high_num=(sorted_bn==highest_thre).nonzero() 70 | percent_limit = high_num.item()/len(bn_weights) 71 | #优化好的模型只有一个极值点，上面代码临时使用以方便调试 72 | # percent_limit = (sorted_bn==highest_thre).nonzero().item()/len(bn_weights) 73 | 74 | print(f'Suggested Threshold should be less than {highest_thre:.4f}.') 75 | print(f'The corresponding prune ratio is {percent_limit:.3f},but you can set higher.') 76 | 77 | 78 | def prune_and_eval(model, sorted_bn, percent=.0): 79 | model_copy = deepcopy(model) 80 | thre_index = int(len(sorted_bn) * percent) 81 | #获得α参数的阈值，小于该值的α参数对应的通道，全部裁剪掉 82 | thre1 = sorted_bn[thre_index] 83 | 84 | print(f'Channels with Gamma value less than {thre1:.6f} are pruned!') 85 | 86 | remain_num = 0 87 | idx_new=dict() 88 | for idx in prune_idx: 89 | 90 | if idx not in shortcut_idx: 91 | 92 | # bn_module = model_copy.module_list[idx][1] 93 | bn_module = model_copy.module_list[idx][1] if type( 94 | model_copy.module_list[idx][1]).__name__ is 'BatchNorm2d' else model_copy.module_list[idx][0] 95 | 96 | mask = obtain_bn_mask(bn_module, thre1) 97 | 98 | #记录剪枝后，每一层卷积层对应的mask 99 | # idx_new[idx]=mask.cpu().numpy() 100 | idx_new[idx]=mask 101 | remain_num += int(mask.sum()) 102 | bn_module.weight.data.mul_(mask) 103 | #bn_module.bias.data.mul_(mask*0.0001) 104 | else: 105 | 106 | bn_module = model_copy.module_list[idx][1] 107 | 108 | 109 | mask=idx_new[shortcut_idx[idx]] 110 | idx_new[idx]=mask 111 | 112 | 113 | remain_num += int(mask.sum()) 114 | bn_module.weight.data.mul_(mask) 115 | 116 | #print(int(mask.sum())) 117 | 118 | with torch.no_grad(): 119 | mAP = eval_model(model_copy)[0][2] 120 | 121 | print(f'Number of channels has been reduced from {len(sorted_bn)} to {remain_num}') 122 | print(f'Prune ratio: {1-remain_num/len(sorted_bn):.3f}') 123 | print(f'mAP of the pruned model is {mAP:.4f}') 124 | 125 | return thre1 126 | 127 | percent = opt.percent 128 | threshold = prune_and_eval(model, sorted_bn, percent) 129 | 130 | 131 | 132 | #**************************************************************** 133 | #虽然上面已经能看到剪枝后的效果，但是没有生成剪枝后的模型结构，因此下面的代码是为了生成新的模型结构并拷贝旧模型参数到新模型 134 | 135 | 136 | 137 | #%% 138 | def obtain_filters_mask(model, thre, CBL_idx, prune_idx): 139 | 140 | pruned = 0 141 | total = 0 142 | num_filters = [] 143 | filters_mask = [] 144 | idx_new=dict() 145 | #CBL_idx存储的是所有带BN的卷积层（YOLO层的前一层卷积层是不带BN的） 146 | for idx in CBL_idx: 147 | # bn_module = model.module_list[idx][1] 148 | bn_module = model.module_list[idx][1] if type(model.module_list[idx][1]).__name__ is 'BatchNorm2d' else \ 149 | model.module_list[idx][0] 150 | if idx in prune_idx: 151 | if idx not in shortcut_idx: 152 | 153 | mask = obtain_bn_mask(bn_module, thre).cpu().numpy() 154 | if type(model.module_list[idx][0]).__name__ is 'BatchNorm2d': 155 | half_num = int(len(mask) / 2) 156 | mask1 = mask[:half_num] 157 | mask2 = mask[half_num:] 158 | remain1 = int(mask1.sum()) 159 | remain2 = int(mask2.sum()) 160 | if remain1 == 0 or remain2 == 0: 161 | print("Channels would be all pruned!") 162 | raise Exception 163 | 164 | idx_new[idx]=mask 165 | remain = int(mask.sum()) 166 | pruned = pruned + mask.shape[0] - remain 167 | 168 | # if remain == 0: 169 | # print("Channels would be all pruned!") 170 | # raise Exception 171 | 172 | # print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 173 | # f'remaining channel: {remain:>4d}') 174 | else: 175 | mask=idx_new[shortcut_idx[idx]] 176 | idx_new[idx]=mask 177 | remain= int(mask.sum()) 178 | pruned = pruned + mask.shape[0] - remain 179 | 180 | if remain == 0: 181 | # print("Channels would be all pruned!") 182 | # raise Exception 183 | max_value = bn_module.weight.data.abs().max() 184 | mask = obtain_bn_mask(bn_module, max_value).cpu().numpy() 185 | remain = int(mask.sum()) 186 | pruned = pruned + mask.shape[0] - remain 187 | 188 | print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 189 | f'remaining channel: {remain:>4d}') 190 | else: 191 | mask = np.ones(bn_module.weight.data.shape) 192 | remain = mask.shape[0] 193 | 194 | total += mask.shape[0] 195 | num_filters.append(remain) 196 | filters_mask.append(mask.copy()) 197 | 198 | #因此，这里求出的prune_ratio,需要裁剪的α参数/cbl_idx中所有的α参数 199 | prune_ratio = pruned / total 200 | print(f'Prune channels: {pruned}\tPrune ratio: {prune_ratio:.3f}') 201 | 202 | return num_filters, filters_mask 203 | 204 | num_filters, filters_mask = obtain_filters_mask(model, threshold, CBL_idx, prune_idx) 205 | 206 | 207 | #CBLidx2mask存储CBL_idx中，每一层BN层对应的mask 208 | CBLidx2mask = {idx: mask for idx, mask in zip(CBL_idx, filters_mask)} 209 | 210 | 211 | pruned_model = prune_model_keep_size2(model, prune_idx, CBL_idx, CBLidx2mask) 212 | print("\nnow prune the model but keep size,(actually add offset of BN beta to next layer), let's see how the mAP goes") 213 | 214 | with torch.no_grad(): 215 | eval_model(pruned_model) 216 | 217 | 218 | 219 | #获得原始模型的module_defs，并修改该defs中的卷积核数量 220 | compact_module_defs = deepcopy(model.module_defs) 221 | for idx, num in zip(CBL_idx, num_filters): 222 | assert compact_module_defs[idx]['type'] == 'convolutional' or compact_module_defs[idx]['type'] == 'convolutional_noconv' 223 | compact_module_defs[idx]['filters'] = str(num) 224 | if compact_module_defs[idx]['type'] == 'convolutional_noconv': 225 | model_def=compact_module_defs[idx-1] #route 226 | assert compact_module_defs[idx-1]['type'] == 'route' 227 | from_layers = [int(s) for s in model_def['layers'].split(',')] 228 | assert compact_module_defs[idx - 1 + from_layers[0]]['type'] == 'convolutional_nobias' 229 | assert compact_module_defs[idx-1 + from_layers[1] if from_layers[1] < 0 else from_layers[1]]['type'] == 'convolutional_nobias' 230 | half_num = int(len(CBLidx2mask[idx]) / 2) 231 | mask1=CBLidx2mask[idx][:half_num] 232 | mask2 = CBLidx2mask[idx][half_num:] 233 | remain1 = int(mask1.sum()) 234 | remain2 = int(mask2.sum()) 235 | compact_module_defs[idx - 1 + from_layers[0]]['filters']=remain1 236 | compact_module_defs[idx-1 + from_layers[1] if from_layers[1] < 0 else from_layers[1]]['filters'] =remain2 237 | 238 | 239 | compact_model = Darknet([model.hyperparams.copy()] + compact_module_defs, (img_size, img_size)).to(device) 240 | compact_nparameters = obtain_num_parameters(compact_model) 241 | 242 | init_weights_from_loose_model(compact_model, pruned_model, CBL_idx, Conv_idx, CBLidx2mask) 243 | 244 | 245 | random_input = torch.rand((1, 3, img_size, img_size)).to(device) 246 | 247 | def obtain_avg_forward_time(input, model, repeat=200): 248 | 249 | model.eval() 250 | start = time.time() 251 | with torch.no_grad(): 252 | for i in range(repeat): 253 | output = model(input)[0] 254 | avg_infer_time = (time.time() - start) / repeat 255 | 256 | return avg_infer_time, output 257 | 258 | print('testing Inference time...') 259 | pruned_forward_time, pruned_output = obtain_avg_forward_time(random_input, pruned_model) 260 | compact_forward_time, compact_output = obtain_avg_forward_time(random_input, compact_model) 261 | 262 | diff = (pruned_output - compact_output).abs().gt(0.001).sum().item() 263 | if diff > 0: 264 | print('Something wrong with the pruned model!') 265 | 266 | # 在测试集上测试剪枝后的模型, 并统计模型的参数数量 267 | print('testing final model') 268 | with torch.no_grad(): 269 | compact_model_metric = eval_model(compact_model) 270 | 271 | 272 | # 比较剪枝前后参数数量的变化、指标性能的变化 273 | metric_table = [ 274 | ["Metric", "Before", "After"], 275 | ["mAP", f'{origin_model_metric[0][2]:.6f}', f'{compact_model_metric[0][2]:.6f}'], 276 | ["Parameters", f"{origin_nparameters}", f"{compact_nparameters}"], 277 | ["Inference", f'{pruned_forward_time:.4f}', f'{compact_forward_time:.4f}'] 278 | ] 279 | print(AsciiTable(metric_table).table) 280 | 281 | 282 | # 生成剪枝后的cfg文件并保存模型 283 | pruned_cfg_name = opt.cfg.replace('/', f'/prune_{percent}_') 284 | pruned_cfg_file = write_cfg(pruned_cfg_name, [model.hyperparams.copy()] + compact_module_defs) 285 | print(f'Config file has been saved: {pruned_cfg_file}') 286 | 287 | compact_model_name = opt.weights.replace('/', f'/prune_{percent}_') 288 | if compact_model_name.endswith('.pt'): 289 | chkpt = {'epoch': -1, 290 | 'best_fitness': None, 291 | 'training_results': None, 292 | 'model': compact_model.state_dict(), 293 | 'optimizer': None} 294 | torch.save(chkpt, compact_model_name) 295 | compact_model_name = compact_model_name.replace('.pt', '.weights') 296 | # save_weights(compact_model, path=compact_model_name) 297 | print(f'Compact model has been saved: {compact_model_name}') 298 | 299 | -------------------------------------------------------------------------------- /slim_prune_yolov5s_8x.py: -------------------------------------------------------------------------------- 1 | from modelsori import * 2 | from utils.utils import * 3 | import numpy as np 4 | from copy import deepcopy 5 | from test import test 6 | from terminaltables import AsciiTable 7 | import time 8 | from utils.prune_utils import * 9 | import argparse 10 | import torchvision 11 | import val 12 | from utils.model_transfer import copy_weight_v6,copy_weight_v6x,copy_weight_v6_reverse,copy_weight_v6x_reverse 13 | 14 | if __name__ == '__main__': 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--cfg', type=str, default='cfg/yolov5s_v6_hand.cfg', help='cfg file path') 17 | parser.add_argument('--data', type=str, default='data/oxfordhand.data', help='*.data file path') 18 | parser.add_argument('--weights', type=str, default='weights/last_v6s.pt', help='sparse model weights') 19 | parser.add_argument('--global_percent', type=float, default=0.6, help='global channel prune percent') 20 | parser.add_argument('--layer_keep', type=float, default=0.01, help='channel keep percent per layer') 21 | parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') 22 | opt = parser.parse_args() 23 | print(opt) 24 | 25 | img_size = opt.img_size 26 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 27 | model = Darknet(opt.cfg, (img_size, img_size)).to(device) 28 | 29 | modelyolov5 = torch.load(opt.weights, map_location=device)['model'].float() # load FP32 model 30 | stride=32.0 31 | if len(modelyolov5.yaml["anchors"]) == 4: 32 | copy_weight_v6x(modelyolov5, model) 33 | stride=64.0 34 | else: 35 | copy_weight_v6(modelyolov5, model) 36 | 37 | if len(modelyolov5.yaml["anchors"]) == 4: 38 | copy_weight_reverse = copy_weight_v6x_reverse 39 | else: 40 | copy_weight_reverse = copy_weight_v6_reverse 41 | 42 | eval_modelv2 = lambda model:val.run(opt.data, 43 | model=model, 44 | batch_size=4, 45 | imgsz=img_size, 46 | plots=False, 47 | stride = stride) 48 | 49 | eval_model = lambda model:test(model=model,cfg=opt.cfg, data=opt.data, batch_size=4, img_size=img_size,stride=stride) 50 | obtain_num_parameters = lambda model:sum([param.nelement() for param in model.parameters()]) 51 | 52 | print("\nlet's test the original model first:") 53 | with torch.no_grad(): 54 | # origin_model_metric = eval_model(model) 55 | copy_weight_reverse(modelyolov5, model) 56 | origin_model_metric = eval_modelv2(modelyolov5) 57 | origin_nparameters = obtain_num_parameters(model) 58 | 59 | CBL_idx, Conv_idx, prune_idx, _, _= parse_module_defs2(model.module_defs) 60 | 61 | 62 | 63 | bn_weights = gather_bn_weights(model.module_list, prune_idx) 64 | 65 | sorted_bn = torch.sort(bn_weights)[0] 66 | sorted_bn, sorted_index = torch.sort(bn_weights) 67 | thresh_index = int(len(bn_weights) * opt.global_percent) 68 | thresh = sorted_bn[thresh_index].cuda() 69 | 70 | print(f'Global Threshold should be less than {thresh:.4f}.') 71 | 72 | #%% 73 | def obtain_filters_mask(model, thre, CBL_idx, prune_idx): 74 | 75 | pruned = 0 76 | total = 0 77 | num_filters = [] 78 | filters_mask = [] 79 | for idx in CBL_idx: 80 | # bn_module = model.module_list[idx][1] 81 | bn_module = model.module_list[idx][1] if type( 82 | model.module_list[idx][1]).__name__ == 'BatchNorm2d' else model.module_list[idx][0] 83 | if idx in prune_idx: 84 | 85 | weight_copy = bn_module.weight.data.abs().clone() 86 | 87 | if model.module_defs[idx][ 'type'] == 'convolutional_noconv': 88 | channels = weight_copy.shape[0] 89 | channels_half=int(channels/2) 90 | weight_copy1=weight_copy[:channels_half] 91 | weight_copy2 = weight_copy[channels_half:] 92 | min_channel_num = int(channels_half * opt.layer_keep) if int(channels_half * opt.layer_keep) > 0 else 1 93 | mask1 = weight_copy1.gt(thresh).float() 94 | mask2 = weight_copy2.gt(thresh).float() 95 | 96 | if int(torch.sum(mask1)) < min_channel_num: 97 | _, sorted_index_weights1 = torch.sort(weight_copy1, descending=True) 98 | mask1[sorted_index_weights1[:min_channel_num]] = 1. 99 | 100 | if int(torch.sum(mask2)) < min_channel_num: 101 | _, sorted_index_weights2 = torch.sort(weight_copy2, descending=True) 102 | mask2[sorted_index_weights2[:min_channel_num]] = 1. 103 | 104 | # regular 105 | mask_cnt1 = int(mask1.sum()) 106 | mask_cnt2 = int(mask2.sum()) 107 | 108 | if mask_cnt1 % 8 != 0: 109 | mask_cnt1 = int((mask_cnt1 // 8 + 1) * 8) 110 | if mask_cnt2 % 8 != 0: 111 | mask_cnt2 = int((mask_cnt2 // 8 + 1) * 8) 112 | 113 | this_layer_sort_bn = bn_module.weight.data.abs().clone() 114 | this_layer_sort_bn1 = this_layer_sort_bn[:channels_half] 115 | this_layer_sort_bn2 = this_layer_sort_bn[channels_half:] 116 | _, sorted_index_weights1 = torch.sort(this_layer_sort_bn1, descending=True) 117 | _, sorted_index_weights2 = torch.sort(this_layer_sort_bn2, descending=True) 118 | mask1[sorted_index_weights1[:mask_cnt1]] = 1. 119 | mask2[sorted_index_weights2[:mask_cnt2]] = 1. 120 | 121 | 122 | remain1 = int(mask1.sum()) 123 | pruned = pruned + mask1.shape[0] - remain1 124 | remain2 = int(mask2.sum()) 125 | pruned = pruned + mask2.shape[0] - remain2 126 | 127 | mask=torch.cat((mask1,mask2)) 128 | remain=remain1+remain2 129 | 130 | print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 131 | f'remaining channel: {remain:>4d}') 132 | else: 133 | 134 | channels = weight_copy.shape[0] # 135 | min_channel_num = int(channels * opt.layer_keep) if int(channels * opt.layer_keep) > 0 else 1 136 | mask = weight_copy.gt(thresh).float() 137 | 138 | if int(torch.sum(mask)) < min_channel_num: 139 | _, sorted_index_weights = torch.sort(weight_copy,descending=True) 140 | mask[sorted_index_weights[:min_channel_num]]=1. 141 | 142 | # regular 143 | mask_cnt = int(mask.sum()) 144 | 145 | if mask_cnt % 8 !=0: 146 | mask_cnt=int((mask_cnt//8+1)*8) 147 | 148 | this_layer_sort_bn = bn_module.weight.data.abs().clone() 149 | _, sorted_index_weights = torch.sort(this_layer_sort_bn,descending=True) 150 | mask[sorted_index_weights[:mask_cnt]]=1. 151 | 152 | remain = int(mask.sum()) 153 | pruned = pruned + mask.shape[0] - remain 154 | 155 | print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t ' 156 | f'remaining channel: {remain:>4d}') 157 | else: 158 | mask = torch.ones(bn_module.weight.data.shape) 159 | remain = mask.shape[0] 160 | 161 | total += mask.shape[0] 162 | num_filters.append(remain) 163 | filters_mask.append(mask.clone()) 164 | 165 | prune_ratio = pruned / total 166 | print(f'Prune channels: {pruned}\tPrune ratio: {prune_ratio:.3f}') 167 | 168 | return num_filters, filters_mask 169 | 170 | num_filters, filters_mask = obtain_filters_mask(model, thresh, CBL_idx, prune_idx) 171 | CBLidx2mask = {idx: mask for idx, mask in zip(CBL_idx, filters_mask)} 172 | CBLidx2filters = {idx: filters for idx, filters in zip(CBL_idx, num_filters)} 173 | 174 | for i in model.module_defs: 175 | if i['type'] == 'shortcut': 176 | i['is_access'] = False 177 | 178 | print('merge the mask of layers connected to shortcut!') 179 | merge_mask_regular(model, CBLidx2mask, CBLidx2filters) 180 | 181 | def prune_and_eval(model, CBL_idx, CBLidx2mask): 182 | model_copy = deepcopy(model) 183 | 184 | for idx in CBL_idx: 185 | # bn_module = model_copy.module_list[idx][1] 186 | bn_module = model_copy.module_list[idx][1] if type( 187 | model_copy.module_list[idx][1]).__name__ == 'BatchNorm2d' else model_copy.module_list[idx][0] 188 | mask = CBLidx2mask[idx].cuda() 189 | bn_module.weight.data.mul_(mask) 190 | 191 | with torch.no_grad(): 192 | # mAP = eval_model(model_copy)[0][2] 193 | copy_weight_reverse(modelyolov5, model_copy) 194 | mAP = eval_modelv2(modelyolov5)[0][2] 195 | 196 | print(f'mask the gamma as zero, mAP of the model is {mAP:.4f}') 197 | 198 | 199 | prune_and_eval(model, CBL_idx, CBLidx2mask) 200 | 201 | 202 | for i in CBLidx2mask: 203 | CBLidx2mask[i] = CBLidx2mask[i].clone().cpu().numpy() 204 | 205 | pruned_model = prune_model_keep_size2(model, prune_idx, CBL_idx, CBLidx2mask) 206 | print("\nnow prune the model but keep size,(actually add offset of BN beta to following layers), let's see how the mAP goes") 207 | 208 | with torch.no_grad(): 209 | # eval_model(pruned_model) 210 | copy_weight_reverse(modelyolov5, pruned_model) 211 | eval_modelv2(modelyolov5) 212 | 213 | 214 | for i in model.module_defs: 215 | if i['type'] == 'shortcut': 216 | i.pop('is_access') 217 | 218 | compact_module_defs = deepcopy(model.module_defs) 219 | for idx in CBL_idx: 220 | assert compact_module_defs[idx]['type'] == 'convolutional' or compact_module_defs[idx][ 221 | 'type'] == 'convolutional_noconv' 222 | num=CBLidx2filters[idx] 223 | compact_module_defs[idx]['filters'] = str(num) 224 | if compact_module_defs[idx]['type'] == 'convolutional_noconv': 225 | model_def = compact_module_defs[idx - 1] # route 226 | assert compact_module_defs[idx - 1]['type'] == 'route' 227 | from_layers = [int(s) for s in model_def['layers'].split(',')] 228 | assert compact_module_defs[idx - 1 + from_layers[0]]['type'] == 'convolutional_nobias' 229 | assert compact_module_defs[idx - 1 + from_layers[1] if from_layers[1] < 0 else from_layers[1]][ 230 | 'type'] == 'convolutional_nobias' 231 | half_num = int(len(CBLidx2mask[idx]) / 2) 232 | mask1 = CBLidx2mask[idx][:half_num] 233 | mask2 = CBLidx2mask[idx][half_num:] 234 | remain1 = int(mask1.sum()) 235 | remain2 = int(mask2.sum()) 236 | compact_module_defs[idx - 1 + from_layers[0]]['filters'] = remain1 237 | compact_module_defs[idx - 1 + from_layers[1] if from_layers[1] < 0 else from_layers[1]]['filters'] = remain2 238 | 239 | 240 | compact_model = Darknet([model.hyperparams.copy()] + compact_module_defs, (img_size, img_size)).to(device) 241 | compact_nparameters = obtain_num_parameters(compact_model) 242 | 243 | init_weights_from_loose_model(compact_model, pruned_model, CBL_idx, Conv_idx, CBLidx2mask) 244 | 245 | 246 | random_input = torch.rand((1, 3, img_size, img_size)).to(device) 247 | 248 | def obtain_avg_forward_time(input, model, repeat=200): 249 | # model.to('cpu').fuse() 250 | # model.module_list.to(device) 251 | model.eval() 252 | start = time.time() 253 | with torch.no_grad(): 254 | for i in range(repeat): 255 | output = model(input)[0] 256 | avg_infer_time = (time.time() - start) / repeat 257 | 258 | return avg_infer_time, output 259 | 260 | print('testing inference time...') 261 | pruned_forward_time, pruned_output = obtain_avg_forward_time(random_input, pruned_model) 262 | compact_forward_time, compact_output = obtain_avg_forward_time(random_input, compact_model) 263 | 264 | diff = (pruned_output - compact_output).abs().gt(0.001).sum().item() 265 | if diff > 0: 266 | print('Something wrong with the pruned model!') 267 | 268 | print('testing the final model...') 269 | with torch.no_grad(): 270 | # compact_model_metric = eval_model(compact_model) 271 | copy_weight_reverse(modelyolov5, compact_model) 272 | compact_model_metric = eval_modelv2(modelyolov5) 273 | 274 | 275 | 276 | metric_table = [ 277 | ["Metric", "Before", "After"], 278 | ["mAP", f'{origin_model_metric[0][2]:.6f}', f'{compact_model_metric[0][2]:.6f}'], 279 | ["Parameters", f"{origin_nparameters}", f"{compact_nparameters}"], 280 | ["Inference", f'{pruned_forward_time:.4f}', f'{compact_forward_time:.4f}'] 281 | ] 282 | print(AsciiTable(metric_table).table) 283 | 284 | 285 | 286 | pruned_cfg_name = opt.cfg.replace('/', f'/prune_{opt.global_percent}_keep_{opt.layer_keep}_8x_') 287 | pruned_cfg_file = write_cfg(pruned_cfg_name, [model.hyperparams.copy()] + compact_module_defs) 288 | print(f'Config file has been saved: {pruned_cfg_file}') 289 | 290 | compact_model_name = opt.weights.replace('/', f'/prune_{opt.global_percent}_keep_{opt.layer_keep}_8x_') 291 | if compact_model_name.endswith('.pt'): 292 | chkpt = {'epoch': -1, 293 | 'best_fitness': None, 294 | 'training_results': None, 295 | 'model': compact_model.state_dict(), 296 | # 'model': compact_model.module_list, #部署调试加载的模型 297 | 'optimizer': None} 298 | torch.save(chkpt, compact_model_name) 299 | compact_model_name = compact_model_name.replace('.pt', '.weights') 300 | # save_weights(compact_model, path=compact_model_name) 301 | print(f'Compact model has been saved: {compact_model_name}') 302 | 303 | -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Model validation metrics 4 | """ 5 | 6 | import math 7 | import warnings 8 | from pathlib import Path 9 | 10 | import matplotlib.pyplot as plt 11 | import numpy as np 12 | import torch 13 | 14 | 15 | def fitness(x): 16 | # Model fitness as a weighted combination of metrics 17 | w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95] 18 | return (x[:, :4] * w).sum(1) 19 | 20 | 21 | def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()): 22 | """ Compute the average precision, given the recall and precision curves. 23 | Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. 24 | # Arguments 25 | tp: True positives (nparray, nx1 or nx10). 26 | conf: Objectness value from 0-1 (nparray). 27 | pred_cls: Predicted object classes (nparray). 28 | target_cls: True object classes (nparray). 29 | plot: Plot precision-recall curve at mAP@0.5 30 | save_dir: Plot save directory 31 | # Returns 32 | The average precision as computed in py-faster-rcnn. 33 | """ 34 | 35 | # Sort by objectness 36 | i = np.argsort(-conf) 37 | tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] 38 | 39 | # Find unique classes 40 | unique_classes = np.unique(target_cls) 41 | nc = unique_classes.shape[0] # number of classes, number of detections 42 | 43 | # Create Precision-Recall curve and compute AP for each class 44 | px, py = np.linspace(0, 1, 1000), [] # for plotting 45 | ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000)) 46 | for ci, c in enumerate(unique_classes): 47 | i = pred_cls == c 48 | n_l = (target_cls == c).sum() # number of labels 49 | n_p = i.sum() # number of predictions 50 | 51 | if n_p == 0 or n_l == 0: 52 | continue 53 | else: 54 | # Accumulate FPs and TPs 55 | fpc = (1 - tp[i]).cumsum(0) 56 | tpc = tp[i].cumsum(0) 57 | 58 | # Recall 59 | recall = tpc / (n_l + 1e-16) # recall curve 60 | r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases 61 | 62 | # Precision 63 | precision = tpc / (tpc + fpc) # precision curve 64 | p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score 65 | 66 | # AP from recall-precision curve 67 | for j in range(tp.shape[1]): 68 | ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j]) 69 | if plot and j == 0: 70 | py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5 71 | 72 | # Compute F1 (harmonic mean of precision and recall) 73 | f1 = 2 * p * r / (p + r + 1e-16) 74 | if plot: 75 | plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names) 76 | plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1') 77 | plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision') 78 | plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall') 79 | 80 | i = f1.mean(0).argmax() # max F1 index 81 | return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32') 82 | 83 | 84 | def compute_ap(recall, precision): 85 | """ Compute the average precision, given the recall and precision curves 86 | # Arguments 87 | recall: The recall curve (list) 88 | precision: The precision curve (list) 89 | # Returns 90 | Average precision, precision curve, recall curve 91 | """ 92 | 93 | # Append sentinel values to beginning and end 94 | mrec = np.concatenate(([0.0], recall, [1.0])) 95 | mpre = np.concatenate(([1.0], precision, [0.0])) 96 | 97 | # Compute the precision envelope 98 | mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) 99 | 100 | # Integrate area under curve 101 | method = 'interp' # methods: 'continuous', 'interp' 102 | if method == 'interp': 103 | x = np.linspace(0, 1, 101) # 101-point interp (COCO) 104 | ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate 105 | else: # 'continuous' 106 | i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes 107 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve 108 | 109 | return ap, mpre, mrec 110 | 111 | 112 | class ConfusionMatrix: 113 | # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix 114 | def __init__(self, nc, conf=0.25, iou_thres=0.45): 115 | self.matrix = np.zeros((nc + 1, nc + 1)) 116 | self.nc = nc # number of classes 117 | self.conf = conf 118 | self.iou_thres = iou_thres 119 | 120 | def process_batch(self, detections, labels): 121 | """ 122 | Return intersection-over-union (Jaccard index) of boxes. 123 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 124 | Arguments: 125 | detections (Array[N, 6]), x1, y1, x2, y2, conf, class 126 | labels (Array[M, 5]), class, x1, y1, x2, y2 127 | Returns: 128 | None, updates confusion matrix accordingly 129 | """ 130 | detections = detections[detections[:, 4] > self.conf] 131 | gt_classes = labels[:, 0].int() 132 | detection_classes = detections[:, 5].int() 133 | iou = box_iou(labels[:, 1:], detections[:, :4]) 134 | 135 | x = torch.where(iou > self.iou_thres) 136 | if x[0].shape[0]: 137 | matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() 138 | if x[0].shape[0] > 1: 139 | matches = matches[matches[:, 2].argsort()[::-1]] 140 | matches = matches[np.unique(matches[:, 1], return_index=True)[1]] 141 | matches = matches[matches[:, 2].argsort()[::-1]] 142 | matches = matches[np.unique(matches[:, 0], return_index=True)[1]] 143 | else: 144 | matches = np.zeros((0, 3)) 145 | 146 | n = matches.shape[0] > 0 147 | m0, m1, _ = matches.transpose().astype(np.int16) 148 | for i, gc in enumerate(gt_classes): 149 | j = m0 == i 150 | if n and sum(j) == 1: 151 | self.matrix[detection_classes[m1[j]], gc] += 1 # correct 152 | else: 153 | self.matrix[self.nc, gc] += 1 # background FP 154 | 155 | if n: 156 | for i, dc in enumerate(detection_classes): 157 | if not any(m1 == i): 158 | self.matrix[dc, self.nc] += 1 # background FN 159 | 160 | def matrix(self): 161 | return self.matrix 162 | 163 | def plot(self, normalize=True, save_dir='', names=()): 164 | try: 165 | import seaborn as sn 166 | 167 | array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-6) if normalize else 1) # normalize columns 168 | array[array < 0.005] = np.nan # don't annotate (would appear as 0.00) 169 | 170 | fig = plt.figure(figsize=(12, 9), tight_layout=True) 171 | sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size 172 | labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels 173 | with warnings.catch_warnings(): 174 | warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered 175 | sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, 176 | xticklabels=names + ['background FP'] if labels else "auto", 177 | yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1)) 178 | fig.axes[0].set_xlabel('True') 179 | fig.axes[0].set_ylabel('Predicted') 180 | fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) 181 | plt.close() 182 | except Exception as e: 183 | print(f'WARNING: ConfusionMatrix plot failure: {e}') 184 | 185 | def print(self): 186 | for i in range(self.nc + 1): 187 | print(' '.join(map(str, self.matrix[i]))) 188 | 189 | 190 | def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7): 191 | # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 192 | box2 = box2.T 193 | 194 | # Get the coordinates of bounding boxes 195 | if x1y1x2y2: # x1, y1, x2, y2 = box1 196 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] 197 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] 198 | else: # transform from xywh to xyxy 199 | b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2 200 | b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2 201 | b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2 202 | b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 203 | 204 | # Intersection area 205 | inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ 206 | (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) 207 | 208 | # Union Area 209 | w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps 210 | w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps 211 | union = w1 * h1 + w2 * h2 - inter + eps 212 | 213 | iou = inter / union 214 | if GIoU or DIoU or CIoU: 215 | cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width 216 | ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height 217 | if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 218 | c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared 219 | rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + 220 | (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared 221 | if DIoU: 222 | return iou - rho2 / c2 # DIoU 223 | elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 224 | v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) 225 | with torch.no_grad(): 226 | alpha = v / (v - iou + (1 + eps)) 227 | return iou - (rho2 / c2 + v * alpha) # CIoU 228 | else: # GIoU https://arxiv.org/pdf/1902.09630.pdf 229 | c_area = cw * ch + eps # convex area 230 | return iou - (c_area - union) / c_area # GIoU 231 | else: 232 | return iou # IoU 233 | 234 | 235 | def box_iou(box1, box2): 236 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 237 | """ 238 | Return intersection-over-union (Jaccard index) of boxes. 239 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 240 | Arguments: 241 | box1 (Tensor[N, 4]) 242 | box2 (Tensor[M, 4]) 243 | Returns: 244 | iou (Tensor[N, M]): the NxM matrix containing the pairwise 245 | IoU values for every element in boxes1 and boxes2 246 | """ 247 | 248 | def box_area(box): 249 | # box = 4xn 250 | return (box[2] - box[0]) * (box[3] - box[1]) 251 | 252 | area1 = box_area(box1.T) 253 | area2 = box_area(box2.T) 254 | 255 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 256 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 257 | return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) 258 | 259 | 260 | def bbox_ioa(box1, box2, eps=1E-7): 261 | """ Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2 262 | box1: np.array of shape(4) 263 | box2: np.array of shape(nx4) 264 | returns: np.array of shape(n) 265 | """ 266 | 267 | box2 = box2.transpose() 268 | 269 | # Get the coordinates of bounding boxes 270 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] 271 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] 272 | 273 | # Intersection area 274 | inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \ 275 | (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0) 276 | 277 | # box2 area 278 | box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + eps 279 | 280 | # Intersection over box2 area 281 | return inter_area / box2_area 282 | 283 | 284 | def wh_iou(wh1, wh2): 285 | # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2 286 | wh1 = wh1[:, None] # [N,1,2] 287 | wh2 = wh2[None] # [1,M,2] 288 | inter = torch.min(wh1, wh2).prod(2) # [N,M] 289 | return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter) 290 | 291 | 292 | # Plots ---------------------------------------------------------------------------------------------------------------- 293 | 294 | def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()): 295 | # Precision-recall curve 296 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) 297 | py = np.stack(py, axis=1) 298 | 299 | if 0 < len(names) < 21: # display per-class legend if < 21 classes 300 | for i, y in enumerate(py.T): 301 | ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}') # plot(recall, precision) 302 | else: 303 | ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision) 304 | 305 | ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean()) 306 | ax.set_xlabel('Recall') 307 | ax.set_ylabel('Precision') 308 | ax.set_xlim(0, 1) 309 | ax.set_ylim(0, 1) 310 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") 311 | fig.savefig(Path(save_dir), dpi=250) 312 | plt.close() 313 | 314 | 315 | def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'): 316 | # Metric-confidence curve 317 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) 318 | 319 | if 0 < len(names) < 21: # display per-class legend if < 21 classes 320 | for i, y in enumerate(py): 321 | ax.plot(px, y, linewidth=1, label=f'{names[i]}') # plot(confidence, metric) 322 | else: 323 | ax.plot(px, py.T, linewidth=1, color='grey') # plot(confidence, metric) 324 | 325 | y = py.mean(0) 326 | ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}') 327 | ax.set_xlabel(xlabel) 328 | ax.set_ylabel(ylabel) 329 | ax.set_xlim(0, 1) 330 | ax.set_ylim(0, 1) 331 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") 332 | fig.savefig(Path(save_dir), dpi=250) 333 | plt.close() 334 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | 4 | from torch.utils.data import DataLoader 5 | 6 | from modelsori import * 7 | from utils.datasets import * 8 | from utils.utils import * 9 | import torchvision 10 | 11 | def box_iouv5(box1, box2): 12 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 13 | """ 14 | Return intersection-over-union (Jaccard index) of boxes. 15 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 16 | Arguments: 17 | box1 (Tensor[N, 4]) 18 | box2 (Tensor[M, 4]) 19 | Returns: 20 | iou (Tensor[N, M]): the NxM matrix containing the pairwise 21 | IoU values for every element in boxes1 and boxes2 22 | """ 23 | 24 | def box_area(box): 25 | # box = 4xn 26 | return (box[2] - box[0]) * (box[3] - box[1]) 27 | 28 | area1 = box_area(box1.T) 29 | area2 = box_area(box2.T) 30 | 31 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 32 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 33 | return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) 34 | 35 | def non_max_suppressionv5(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): 36 | """Performs Non-Maximum Suppression (NMS) on inference results 37 | 38 | Returns: 39 | detections with shape: nx6 (x1, y1, x2, y2, conf, cls) 40 | """ 41 | if prediction.dtype is torch.float16: 42 | prediction = prediction.float() # to FP32 43 | 44 | nc = prediction[0].shape[1] - 5 # number of classes 45 | xc = prediction[..., 4] > conf_thres # candidates 46 | 47 | # Settings 48 | min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height 49 | max_det = 300 # maximum number of detections per image 50 | time_limit = 10.0 # seconds to quit after 51 | redundant = True # require redundant detections 52 | multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) 53 | 54 | t = time.time() 55 | output = [None] * prediction.shape[0] 56 | for xi, x in enumerate(prediction): # image index, image inference 57 | # Apply constraints 58 | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height 59 | x = x[xc[xi]] # confidence 60 | 61 | # If none remain process next image 62 | if not x.shape[0]: 63 | continue 64 | 65 | # Compute conf 66 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 67 | 68 | # Box (center x, center y, width, height) to (x1, y1, x2, y2) 69 | box = xywh2xyxy(x[:, :4]) 70 | 71 | # Detections matrix nx6 (xyxy, conf, cls) 72 | if multi_label: 73 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 74 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 75 | else: # best class only 76 | conf, j = x[:, 5:].max(1, keepdim=True) 77 | x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] 78 | 79 | # Filter by class 80 | if classes: 81 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 82 | 83 | # Apply finite constraint 84 | # if not torch.isfinite(x).all(): 85 | # x = x[torch.isfinite(x).all(1)] 86 | 87 | # If none remain process next image 88 | n = x.shape[0] # number of boxes 89 | if not n: 90 | continue 91 | 92 | # Sort by confidence 93 | # x = x[x[:, 4].argsort(descending=True)] 94 | 95 | # Batched NMS 96 | c = x[:, 5:6] * (0 if agnostic else max_wh) # classes 97 | boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores 98 | i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) 99 | if i.shape[0] > max_det: # limit detections 100 | i = i[:max_det] 101 | if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) 102 | try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) 103 | iou = box_iouv5(boxes[i], boxes) > iou_thres # iou matrix 104 | weights = iou * scores[None] # box weights 105 | x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes 106 | if redundant: 107 | i = i[iou.sum(1) > 1] # require redundancy 108 | except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 109 | print(x, i, x.shape, i.shape) 110 | pass 111 | 112 | output[xi] = x[i] 113 | if (time.time() - t) > time_limit: 114 | break # time limit exceeded 115 | 116 | return output 117 | 118 | def test(cfg, 119 | data, 120 | weights=None, 121 | batch_size=16, 122 | img_size=416, 123 | iou_thres=0.5, 124 | conf_thres=0.001, 125 | nms_thres=0.5, 126 | stride=32.0, 127 | save_json=False, 128 | model=None): 129 | 130 | # Initialize/load model and set device 131 | if model is None: 132 | device = torch_utils.select_device(opt.device) 133 | verbose = True 134 | 135 | # Initialize model 136 | model = Darknet(cfg, img_size).to(device) 137 | 138 | # Load weights 139 | attempt_download(weights) 140 | if weights.endswith('.pt'): # pytorch format 141 | model.load_state_dict(torch.load(weights, map_location=device)['model']) 142 | else: # darknet format 143 | _ = load_darknet_weights(model, weights) 144 | 145 | if torch.cuda.device_count() > 1: 146 | model = nn.DataParallel(model) 147 | else: 148 | device = next(model.parameters()).device # get model device 149 | verbose = False 150 | 151 | # Configure run 152 | data = parse_data_cfg(data) 153 | nc = int(data['classes']) # number of classes 154 | test_path = data['valid'] # path to test images 155 | names = load_classes(data['names']) # class names 156 | 157 | # Dataloader 158 | dataset = LoadImagesAndLabels(test_path, img_size, batch_size,stride=stride) 159 | dataloader = DataLoader(dataset, 160 | batch_size=batch_size, 161 | num_workers=min([os.cpu_count(), batch_size, 16]), 162 | pin_memory=True, 163 | collate_fn=dataset.collate_fn) 164 | 165 | seen = 0 166 | model.eval() 167 | coco91class = coco80_to_coco91_class() 168 | s = ('%20s' + '%10s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP', 'F1') 169 | p, r, f1, mp, mr, map, mf1 = 0., 0., 0., 0., 0., 0., 0. 170 | loss = torch.zeros(3) 171 | jdict, stats, ap, ap_class = [], [], [], [] 172 | for batch_i, (imgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)): 173 | targets = targets.to(device) 174 | imgs = imgs.to(device) 175 | _, _, height, width = imgs.shape # batch size, channels, height, width 176 | 177 | # Plot images with bounding boxes 178 | if batch_i == 0 and not os.path.exists('test_batch0.jpg'): 179 | plot_images(imgs=imgs, targets=targets, paths=paths, fname='test_batch0.jpg') 180 | 181 | # Run model 182 | inf_out, train_out = model(imgs) # inference and training outputs 183 | 184 | # Compute loss 185 | if hasattr(model, 'hyp'): # if model has loss hyperparameters 186 | loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls 187 | 188 | # Run NMS 189 | # output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres) 190 | output = non_max_suppressionv5(inf_out,conf_thres=conf_thres, iou_thres=nms_thres, classes=None,agnostic=False) 191 | 192 | # Statistics per image 193 | for si, pred in enumerate(output): 194 | labels = targets[targets[:, 0] == si, 1:] 195 | nl = len(labels) 196 | tcls = labels[:, 0].tolist() if nl else [] # target class 197 | seen += 1 198 | 199 | if pred is None: 200 | if nl: 201 | stats.append(([], torch.Tensor(), torch.Tensor(), tcls)) 202 | continue 203 | 204 | # Append to text file 205 | # with open('test.txt', 'a') as file: 206 | # [file.write('%11.5g' * 7 % tuple(x) + '\n') for x in pred] 207 | 208 | # Append to pycocotools JSON dictionary 209 | if save_json: 210 | # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... 211 | image_id = int(Path(paths[si]).stem.split('_')[-1]) 212 | box = pred[:, :4].clone() # xyxy 213 | scale_coords(imgs[si].shape[1:], box, shapes[si]) # to original shape 214 | box = xyxy2xywh(box) # xywh 215 | box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner 216 | for di, d in enumerate(pred): 217 | jdict.append({'image_id': image_id, 218 | 'category_id': coco91class[int(d[6])], 219 | 'bbox': [floatn(x, 3) for x in box[di]], 220 | 'score': floatn(d[4], 5)}) 221 | 222 | # Clip boxes to image bounds 223 | clip_coords(pred, (height, width)) 224 | 225 | # Assign all predictions as incorrect 226 | correct = [0] * len(pred) 227 | if nl: 228 | detected = [] 229 | tcls_tensor = labels[:, 0] 230 | 231 | # target boxes 232 | tbox = xywh2xyxy(labels[:, 1:5]) 233 | tbox[:, [0, 2]] *= width 234 | tbox[:, [1, 3]] *= height 235 | 236 | # Search for correct predictions 237 | for i, (*pbox, pconf, pcls) in enumerate(pred): 238 | # for i, (*pbox, pconf, pcls_conf, pcls) in enumerate(pred): 239 | 240 | # Break if all targets already located in image 241 | if len(detected) == nl: 242 | break 243 | 244 | # Continue if predicted class not among image classes 245 | if pcls.item() not in tcls: 246 | continue 247 | 248 | # Best iou, index between pred and targets 249 | m = (pcls == tcls_tensor).nonzero(as_tuple=False).view(-1) 250 | iou, bi = bbox_iou(pbox, tbox[m]).max(0) 251 | 252 | # If iou > threshold and class is correct mark as correct 253 | if iou > iou_thres and m[bi] not in detected: # and pcls == tcls[bi]: 254 | correct[i] = 1 255 | detected.append(m[bi]) 256 | 257 | # Append statistics (correct, conf, pcls, tcls) 258 | # stats.append((correct, pred[:, 4].cpu(), pred[:, 6].cpu(), tcls)) 259 | stats.append((correct, pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) 260 | 261 | # Compute statistics 262 | stats = [np.concatenate(x, 0) for x in list(zip(*stats))] # to numpy 263 | if len(stats): 264 | p, r, ap, f1, ap_class = ap_per_class(*stats) 265 | mp, mr, map, mf1 = p.mean(), r.mean(), ap.mean(), f1.mean() 266 | nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class 267 | else: 268 | nt = torch.zeros(1) 269 | 270 | # Print results 271 | pf = '%20s' + '%10.3g' * 6 # print format 272 | print(pf % ('all', seen, nt.sum(), mp, mr, map, mf1)) 273 | 274 | # Print results per class 275 | if verbose and nc > 1 and len(stats): 276 | for i, c in enumerate(ap_class): 277 | print(pf % (names[c], seen, nt[c], p[i], r[i], ap[i], f1[i])) 278 | 279 | # Save JSON 280 | if save_json and map and len(jdict): 281 | try: 282 | imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataset.img_files] 283 | with open('results.json', 'w') as file: 284 | json.dump(jdict, file) 285 | 286 | from pycocotools.coco import COCO 287 | from pycocotools.cocoeval import COCOeval 288 | 289 | # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb 290 | cocoGt = COCO('../coco/annotations/instances_val2014.json') # initialize COCO ground truth api 291 | cocoDt = cocoGt.loadRes('results.json') # initialize COCO pred api 292 | 293 | cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') 294 | cocoEval.params.imgIds = imgIds # [:32] # only evaluate these images 295 | cocoEval.evaluate() 296 | cocoEval.accumulate() 297 | cocoEval.summarize() 298 | map = cocoEval.stats[1] # update mAP to pycocotools mAP 299 | except: 300 | print('WARNING: missing dependency pycocotools from requirements.txt. Can not compute official COCO mAP.') 301 | 302 | # Return results 303 | maps = np.zeros(nc) + map 304 | for i, c in enumerate(ap_class): 305 | maps[c] = ap[i] 306 | return (mp, mr, map, mf1, *(loss / len(dataloader)).tolist()), maps 307 | 308 | 309 | if __name__ == '__main__': 310 | parser = argparse.ArgumentParser(prog='test.py') 311 | parser.add_argument('--cfg', type=str, default='cfg/yolov5s.cfg', help='cfg file path') 312 | parser.add_argument('--data', type=str, default='data/coco.data', help='coco.data file path') 313 | parser.add_argument('--weights', type=str, default='weights/yolov3-spp.weights', help='path to weights file') 314 | parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch') 315 | parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)') 316 | parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected') 317 | parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') 318 | parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression') 319 | parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') 320 | parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu') 321 | opt = parser.parse_args() 322 | print(opt) 323 | 324 | with torch.no_grad(): 325 | test(opt.cfg, 326 | opt.data, 327 | opt.weights, 328 | opt.batch_size, 329 | opt.img_size, 330 | opt.iou_thres, 331 | opt.conf_thres, 332 | opt.nms_thres, 333 | opt.save_json) 334 | -------------------------------------------------------------------------------- /models/yolo.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | YOLO-specific modules 4 | 5 | Usage: 6 | $ python path/to/models/yolo.py --cfg yolov5s.yaml 7 | """ 8 | 9 | import argparse 10 | import sys 11 | from copy import deepcopy 12 | from pathlib import Path 13 | 14 | FILE = Path(__file__).resolve() 15 | ROOT = FILE.parents[1] # YOLOv5 root directory 16 | if str(ROOT) not in sys.path: 17 | sys.path.append(str(ROOT)) # add ROOT to PATH 18 | # ROOT = ROOT.relative_to(Path.cwd()) # relative 19 | 20 | from models.common import * 21 | from models.experimental import * 22 | from utils.autoanchor import check_anchor_order 23 | from utils.general import check_yaml, make_divisible, print_args, set_logging 24 | from utils.plots import feature_visualization 25 | from utils.torch_utils import copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, \ 26 | select_device, time_sync 27 | 28 | try: 29 | import thop # for FLOPs computation 30 | except ImportError: 31 | thop = None 32 | 33 | LOGGER = logging.getLogger(__name__) 34 | 35 | 36 | class Detect(nn.Module): 37 | stride = None # strides computed during build 38 | onnx_dynamic = False # ONNX export parameter 39 | 40 | def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer 41 | super().__init__() 42 | self.nc = nc # number of classes 43 | self.no = nc + 5 # number of outputs per anchor 44 | self.nl = len(anchors) # number of detection layers 45 | self.na = len(anchors[0]) // 2 # number of anchors 46 | self.grid = [torch.zeros(1)] * self.nl # init grid 47 | self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid 48 | self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2) 49 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 50 | self.inplace = inplace # use in-place ops (e.g. slice assignment) 51 | 52 | def forward(self, x): 53 | z = [] # inference output 54 | for i in range(self.nl): 55 | x[i] = self.m[i](x[i]) # conv 56 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 57 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 58 | 59 | if not self.training: # inference 60 | if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic: 61 | self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i) 62 | 63 | y = x[i].sigmoid() 64 | if self.inplace: 65 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy 66 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 67 | else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953 68 | xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy 69 | wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 70 | y = torch.cat((xy, wh, y[..., 4:]), -1) 71 | z.append(y.view(bs, -1, self.no)) 72 | 73 | return x if self.training else (torch.cat(z, 1), x) 74 | 75 | def _make_grid(self, nx=20, ny=20, i=0): 76 | d = self.anchors[i].device 77 | yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)]) 78 | grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float() 79 | anchor_grid = (self.anchors[i].clone() * self.stride[i]) \ 80 | .view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float() 81 | return grid, anchor_grid 82 | 83 | 84 | class Model(nn.Module): 85 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes 86 | super().__init__() 87 | if isinstance(cfg, dict): 88 | self.yaml = cfg # model dict 89 | else: # is *.yaml 90 | import yaml # for torch hub 91 | self.yaml_file = Path(cfg).name 92 | with open(cfg, errors='ignore') as f: 93 | self.yaml = yaml.safe_load(f) # model dict 94 | 95 | # Define model 96 | ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels 97 | if nc and nc != self.yaml['nc']: 98 | LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}") 99 | self.yaml['nc'] = nc # override yaml value 100 | if anchors: 101 | LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}') 102 | self.yaml['anchors'] = round(anchors) # override yaml value 103 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist 104 | self.names = [str(i) for i in range(self.yaml['nc'])] # default names 105 | self.inplace = self.yaml.get('inplace', True) 106 | 107 | # Build strides, anchors 108 | m = self.model[-1] # Detect() 109 | if isinstance(m, Detect): 110 | s = 256 # 2x min stride 111 | m.inplace = self.inplace 112 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 113 | m.anchors /= m.stride.view(-1, 1, 1) 114 | check_anchor_order(m) 115 | self.stride = m.stride 116 | self._initialize_biases() # only run once 117 | 118 | # Init weights, biases 119 | initialize_weights(self) 120 | self.info() 121 | LOGGER.info('') 122 | 123 | def forward(self, x, augment=False, profile=False, visualize=False): 124 | if augment: 125 | return self._forward_augment(x) # augmented inference, None 126 | return self._forward_once(x, profile, visualize) # single-scale inference, train 127 | 128 | def _forward_augment(self, x): 129 | img_size = x.shape[-2:] # height, width 130 | s = [1, 0.83, 0.67] # scales 131 | f = [None, 3, None] # flips (2-ud, 3-lr) 132 | y = [] # outputs 133 | for si, fi in zip(s, f): 134 | xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max())) 135 | yi = self._forward_once(xi)[0] # forward 136 | # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 137 | yi = self._descale_pred(yi, fi, si, img_size) 138 | y.append(yi) 139 | y = self._clip_augmented(y) # clip augmented tails 140 | return torch.cat(y, 1), None # augmented inference, train 141 | 142 | def _forward_once(self, x, profile=False, visualize=False): 143 | y, dt = [], [] # outputs 144 | for m in self.model: 145 | if m.f != -1: # if not from previous layer 146 | 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 147 | if profile: 148 | self._profile_one_layer(m, x, dt) 149 | x = m(x) # run 150 | y.append(x if m.i in self.save else None) # save output 151 | if visualize: 152 | feature_visualization(x, m.type, m.i, save_dir=visualize) 153 | return x 154 | 155 | def _descale_pred(self, p, flips, scale, img_size): 156 | # de-scale predictions following augmented inference (inverse operation) 157 | if self.inplace: 158 | p[..., :4] /= scale # de-scale 159 | if flips == 2: 160 | p[..., 1] = img_size[0] - p[..., 1] # de-flip ud 161 | elif flips == 3: 162 | p[..., 0] = img_size[1] - p[..., 0] # de-flip lr 163 | else: 164 | x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale # de-scale 165 | if flips == 2: 166 | y = img_size[0] - y # de-flip ud 167 | elif flips == 3: 168 | x = img_size[1] - x # de-flip lr 169 | p = torch.cat((x, y, wh, p[..., 4:]), -1) 170 | return p 171 | 172 | def _clip_augmented(self, y): 173 | # Clip YOLOv5 augmented inference tails 174 | nl = self.model[-1].nl # number of detection layers (P3-P5) 175 | g = sum(4 ** x for x in range(nl)) # grid points 176 | e = 1 # exclude layer count 177 | i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e)) # indices 178 | y[0] = y[0][:, :-i] # large 179 | i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # indices 180 | y[-1] = y[-1][:, i:] # small 181 | return y 182 | 183 | def _profile_one_layer(self, m, x, dt): 184 | c = isinstance(m, Detect) # is final layer, copy input as inplace fix 185 | o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs 186 | t = time_sync() 187 | for _ in range(10): 188 | m(x.copy() if c else x) 189 | dt.append((time_sync() - t) * 100) 190 | if m == self.model[0]: 191 | LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s} {'module'}") 192 | LOGGER.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}') 193 | if c: 194 | LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total") 195 | 196 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 197 | # https://arxiv.org/abs/1708.02002 section 3.3 198 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 199 | m = self.model[-1] # Detect() module 200 | for mi, s in zip(m.m, m.stride): # from 201 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 202 | b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 203 | b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 204 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 205 | 206 | def _print_biases(self): 207 | m = self.model[-1] # Detect() module 208 | for mi in m.m: # from 209 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 210 | LOGGER.info( 211 | ('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 212 | 213 | # def _print_weights(self): 214 | # for m in self.model.modules(): 215 | # if type(m) is Bottleneck: 216 | # LOGGER.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 217 | 218 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 219 | LOGGER.info('Fusing layers... ') 220 | for m in self.model.modules(): 221 | if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'): 222 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 223 | delattr(m, 'bn') # remove batchnorm 224 | m.forward = m.forward_fuse # update forward 225 | self.info() 226 | return self 227 | 228 | def autoshape(self): # add AutoShape module 229 | LOGGER.info('Adding AutoShape... ') 230 | m = AutoShape(self) # wrap model 231 | copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes 232 | return m 233 | 234 | def info(self, verbose=False, img_size=640): # print model information 235 | model_info(self, verbose, img_size) 236 | 237 | def _apply(self, fn): 238 | # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers 239 | self = super()._apply(fn) 240 | m = self.model[-1] # Detect() 241 | if isinstance(m, Detect): 242 | m.stride = fn(m.stride) 243 | m.grid = list(map(fn, m.grid)) 244 | if isinstance(m.anchor_grid, list): 245 | m.anchor_grid = list(map(fn, m.anchor_grid)) 246 | return self 247 | 248 | 249 | def parse_model(d, ch): # model_dict, input_channels(3) 250 | LOGGER.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 251 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 252 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 253 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 254 | 255 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 256 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 257 | m = eval(m) if isinstance(m, str) else m # eval strings 258 | for j, a in enumerate(args): 259 | try: 260 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 261 | except NameError: 262 | pass 263 | 264 | n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain 265 | if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, 266 | BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]: 267 | c1, c2 = ch[f], args[0] 268 | if c2 != no: # if not output 269 | c2 = make_divisible(c2 * gw, 8) 270 | 271 | args = [c1, c2, *args[1:]] 272 | if m in [BottleneckCSP, C3, C3TR, C3Ghost]: 273 | args.insert(2, n) # number of repeats 274 | n = 1 275 | elif m is nn.BatchNorm2d: 276 | args = [ch[f]] 277 | elif m is Concat: 278 | c2 = sum([ch[x] for x in f]) 279 | elif m is Detect: 280 | args.append([ch[x] for x in f]) 281 | if isinstance(args[1], int): # number of anchors 282 | args[1] = [list(range(args[1] * 2))] * len(f) 283 | elif m is Contract: 284 | c2 = ch[f] * args[0] ** 2 285 | elif m is Expand: 286 | c2 = ch[f] // args[0] ** 2 287 | else: 288 | c2 = ch[f] 289 | 290 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 291 | t = str(m)[8:-2].replace('__main__.', '') # module type 292 | np = sum([x.numel() for x in m_.parameters()]) # number params 293 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 294 | LOGGER.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n_, np, t, args)) # print 295 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 296 | layers.append(m_) 297 | if i == 0: 298 | ch = [] 299 | ch.append(c2) 300 | return nn.Sequential(*layers), sorted(save) 301 | 302 | 303 | if __name__ == '__main__': 304 | parser = argparse.ArgumentParser() 305 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 306 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 307 | parser.add_argument('--profile', action='store_true', help='profile model speed') 308 | opt = parser.parse_args() 309 | opt.cfg = check_yaml(opt.cfg) # check YAML 310 | print_args(FILE.stem, opt) 311 | set_logging() 312 | device = select_device(opt.device) 313 | 314 | # Create model 315 | model = Model(opt.cfg).to(device) 316 | model.train() 317 | 318 | # Profile 319 | if opt.profile: 320 | img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 321 | y = model(img, profile=True) 322 | 323 | # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) 324 | # from torch.utils.tensorboard import SummaryWriter 325 | # tb_writer = SummaryWriter('.') 326 | # LOGGER.info("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/") 327 | # tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph 328 | -------------------------------------------------------------------------------- /test_yolov5s.py: -------------------------------------------------------------------------------- 1 | from modelsori import * 2 | from utils.utils import * 3 | import numpy as np 4 | from copy import deepcopy 5 | from test import test 6 | from terminaltables import AsciiTable 7 | import time 8 | from utils.prune_utils import * 9 | import argparse 10 | 11 | from models.yolo import Model 12 | 13 | import torchvision 14 | 15 | def letterboxv5(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): 16 | # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 17 | shape = img.shape[:2] # current shape [height, width] 18 | if isinstance(new_shape, int): 19 | new_shape = (new_shape, new_shape) 20 | 21 | # Scale ratio (new / old) 22 | r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) 23 | if not scaleup: # only scale down, do not scale up (for better test mAP) 24 | r = min(r, 1.0) 25 | 26 | # Compute padding 27 | ratio = r, r # width, height ratios 28 | new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) 29 | dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding 30 | if auto: # minimum rectangle 31 | # dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding 32 | dw, dh = np.mod(dw, 32), np.mod(dh, 32) # wh padding 33 | elif scaleFill: # stretch 34 | dw, dh = 0.0, 0.0 35 | new_unpad = (new_shape[1], new_shape[0]) 36 | ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios 37 | 38 | dw /= 2 # divide padding into 2 sides 39 | dh /= 2 40 | 41 | if shape[::-1] != new_unpad: # resize 42 | img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) 43 | top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) 44 | left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) 45 | img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border 46 | return img, ratio, (dw, dh) 47 | 48 | def box_iou(box1, box2): 49 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 50 | """ 51 | Return intersection-over-union (Jaccard index) of boxes. 52 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 53 | Arguments: 54 | box1 (Tensor[N, 4]) 55 | box2 (Tensor[M, 4]) 56 | Returns: 57 | iou (Tensor[N, M]): the NxM matrix containing the pairwise 58 | IoU values for every element in boxes1 and boxes2 59 | """ 60 | 61 | def box_area(box): 62 | # box = 4xn 63 | return (box[2] - box[0]) * (box[3] - box[1]) 64 | 65 | area1 = box_area(box1.T) 66 | area2 = box_area(box2.T) 67 | 68 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 69 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 70 | return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) 71 | 72 | def non_max_suppressionv5(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): 73 | """Performs Non-Maximum Suppression (NMS) on inference results 74 | 75 | Returns: 76 | detections with shape: nx6 (x1, y1, x2, y2, conf, cls) 77 | """ 78 | if prediction.dtype is torch.float16: 79 | prediction = prediction.float() # to FP32 80 | 81 | nc = prediction[0].shape[1] - 5 # number of classes 82 | xc = prediction[..., 4] > conf_thres # candidates 83 | 84 | # Settings 85 | min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height 86 | max_det = 300 # maximum number of detections per image 87 | time_limit = 10.0 # seconds to quit after 88 | redundant = True # require redundant detections 89 | multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) 90 | 91 | t = time.time() 92 | output = [None] * prediction.shape[0] 93 | for xi, x in enumerate(prediction): # image index, image inference 94 | # Apply constraints 95 | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height 96 | x = x[xc[xi]] # confidence 97 | 98 | # If none remain process next image 99 | if not x.shape[0]: 100 | continue 101 | 102 | # Compute conf 103 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 104 | 105 | # Box (center x, center y, width, height) to (x1, y1, x2, y2) 106 | box = xywh2xyxy(x[:, :4]) 107 | 108 | # Detections matrix nx6 (xyxy, conf, cls) 109 | if multi_label: 110 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 111 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 112 | else: # best class only 113 | conf, j = x[:, 5:].max(1, keepdim=True) 114 | x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] 115 | 116 | # Filter by class 117 | if classes: 118 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 119 | 120 | # Apply finite constraint 121 | # if not torch.isfinite(x).all(): 122 | # x = x[torch.isfinite(x).all(1)] 123 | 124 | # If none remain process next image 125 | n = x.shape[0] # number of boxes 126 | if not n: 127 | continue 128 | 129 | # Sort by confidence 130 | # x = x[x[:, 4].argsort(descending=True)] 131 | 132 | # Batched NMS 133 | c = x[:, 5:6] * (0 if agnostic else max_wh) # classes 134 | boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores 135 | i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) 136 | if i.shape[0] > max_det: # limit detections 137 | i = i[:max_det] 138 | if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) 139 | try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) 140 | iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix 141 | weights = iou * scores[None] # box weights 142 | x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes 143 | if redundant: 144 | i = i[iou.sum(1) > 1] # require redundancy 145 | except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 146 | print(x, i, x.shape, i.shape) 147 | pass 148 | 149 | output[xi] = x[i] 150 | if (time.time() - t) > time_limit: 151 | break # time limit exceeded 152 | 153 | return output 154 | 155 | def plot_one_box(x, img, color=None, label=None, line_thickness=None): 156 | # Plots one bounding box on image img 157 | tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness 158 | color = color or [random.randint(0, 255) for _ in range(3)] 159 | c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) 160 | cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) 161 | if label: 162 | tf = max(tl - 1, 1) # font thickness 163 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] 164 | c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 165 | cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled 166 | cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) 167 | 168 | def copy_conv(conv_src,conv_dst): 169 | conv_dst[0] = conv_src.conv 170 | conv_dst[1] = conv_src.bn 171 | conv_dst[2] = conv_src.act 172 | 173 | def copy_weight_v4(modelyolov5,model): 174 | focus = list(modelyolov5.model.children())[0] 175 | copy_conv(focus.conv, model.module_list[1]) 176 | conv1 = list(modelyolov5.model.children())[1] 177 | copy_conv(conv1, model.module_list[2]) 178 | cspnet1 = list(modelyolov5.model.children())[2] 179 | copy_conv(cspnet1.cv2, model.module_list[3]) 180 | copy_conv(cspnet1.cv1, model.module_list[5]) 181 | copy_conv(cspnet1.m[0].cv1, model.module_list[6]) 182 | copy_conv(cspnet1.m[0].cv2, model.module_list[7]) 183 | copy_conv(cspnet1.cv3, model.module_list[10]) 184 | conv2 = list(modelyolov5.model.children())[3] 185 | copy_conv(conv2, model.module_list[11]) 186 | cspnet2 = list(modelyolov5.model.children())[4] 187 | copy_conv(cspnet2.cv2, model.module_list[12]) 188 | copy_conv(cspnet2.cv1, model.module_list[14]) 189 | copy_conv(cspnet2.m[0].cv1, model.module_list[15]) 190 | copy_conv(cspnet2.m[0].cv2, model.module_list[16]) 191 | copy_conv(cspnet2.m[1].cv1, model.module_list[18]) 192 | copy_conv(cspnet2.m[1].cv2, model.module_list[19]) 193 | copy_conv(cspnet2.m[2].cv1, model.module_list[21]) 194 | copy_conv(cspnet2.m[2].cv2, model.module_list[22]) 195 | copy_conv(cspnet2.cv3, model.module_list[25]) 196 | conv3 = list(modelyolov5.model.children())[5] 197 | copy_conv(conv3, model.module_list[26]) 198 | cspnet3 = list(modelyolov5.model.children())[6] 199 | copy_conv(cspnet3.cv2, model.module_list[27]) 200 | copy_conv(cspnet3.cv1, model.module_list[29]) 201 | copy_conv(cspnet3.m[0].cv1, model.module_list[30]) 202 | copy_conv(cspnet3.m[0].cv2, model.module_list[31]) 203 | copy_conv(cspnet3.m[1].cv1, model.module_list[33]) 204 | copy_conv(cspnet3.m[1].cv2, model.module_list[34]) 205 | copy_conv(cspnet3.m[2].cv1, model.module_list[36]) 206 | copy_conv(cspnet3.m[2].cv2, model.module_list[37]) 207 | copy_conv(cspnet3.cv3, model.module_list[40]) 208 | conv4 = list(modelyolov5.model.children())[7] 209 | copy_conv(conv4, model.module_list[41]) 210 | spp = list(modelyolov5.model.children())[8] 211 | copy_conv(spp.cv1, model.module_list[42]) 212 | model.module_list[43] = spp.m[0] 213 | model.module_list[45] = spp.m[1] 214 | model.module_list[47] = spp.m[2] 215 | copy_conv(spp.cv2, model.module_list[49]) 216 | cspnet4 = list(modelyolov5.model.children())[9] 217 | copy_conv(cspnet4.cv2, model.module_list[50]) 218 | copy_conv(cspnet4.cv1, model.module_list[52]) 219 | copy_conv(cspnet4.m[0].cv1, model.module_list[53]) 220 | copy_conv(cspnet4.m[0].cv2, model.module_list[54]) 221 | copy_conv(cspnet4.cv3, model.module_list[56]) 222 | conv5 = list(modelyolov5.model.children())[10] 223 | copy_conv(conv5, model.module_list[57]) 224 | upsample1 = list(modelyolov5.model.children())[11] 225 | model.module_list[58] = upsample1 226 | cspnet5 = list(modelyolov5.model.children())[13] 227 | copy_conv(cspnet5.cv2, model.module_list[60]) 228 | copy_conv(cspnet5.cv1, model.module_list[62]) 229 | copy_conv(cspnet5.m[0].cv1, model.module_list[63]) 230 | copy_conv(cspnet5.m[0].cv2, model.module_list[64]) 231 | copy_conv(cspnet5.cv3, model.module_list[66]) 232 | conv6 = list(modelyolov5.model.children())[14] 233 | copy_conv(conv6, model.module_list[67]) 234 | upsample2 = list(modelyolov5.model.children())[15] 235 | model.module_list[68] = upsample2 236 | cspnet6 = list(modelyolov5.model.children())[17] 237 | copy_conv(cspnet6.cv2, model.module_list[70]) 238 | copy_conv(cspnet6.cv1, model.module_list[72]) 239 | copy_conv(cspnet6.m[0].cv1, model.module_list[73]) 240 | copy_conv(cspnet6.m[0].cv2, model.module_list[74]) 241 | copy_conv(cspnet6.cv3, model.module_list[76]) 242 | conv7 = list(modelyolov5.model.children())[18] 243 | copy_conv(conv7, model.module_list[80]) 244 | cspnet7 = list(modelyolov5.model.children())[20] 245 | copy_conv(cspnet7.cv2, model.module_list[82]) 246 | copy_conv(cspnet7.cv1, model.module_list[84]) 247 | copy_conv(cspnet7.m[0].cv1, model.module_list[85]) 248 | copy_conv(cspnet7.m[0].cv2, model.module_list[86]) 249 | copy_conv(cspnet7.cv3, model.module_list[88]) 250 | conv8 = list(modelyolov5.model.children())[21] 251 | copy_conv(conv8, model.module_list[92]) 252 | cspnet8 = list(modelyolov5.model.children())[23] 253 | copy_conv(cspnet8.cv2, model.module_list[94]) 254 | copy_conv(cspnet8.cv1, model.module_list[96]) 255 | copy_conv(cspnet8.m[0].cv1, model.module_list[97]) 256 | copy_conv(cspnet8.m[0].cv2, model.module_list[98]) 257 | copy_conv(cspnet8.cv3, model.module_list[100]) 258 | detect = list(modelyolov5.model.children())[24] 259 | model.module_list[77][0] = detect.m[0] 260 | model.module_list[89][0] = detect.m[1] 261 | model.module_list[101][0] = detect.m[2] 262 | 263 | def initialize_weights(model): 264 | for m in model.modules(): 265 | t = type(m) 266 | if t is nn.Conv2d: 267 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 268 | elif t is nn.BatchNorm2d: 269 | m.eps = 1e-3 270 | m.momentum = 0.03 271 | elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 272 | m.inplace = True 273 | 274 | if __name__ == '__main__': 275 | parser = argparse.ArgumentParser() 276 | parser.add_argument('--cfg', type=str, default='cfg/yolov5s_v4.cfg', help='cfg file path') 277 | parser.add_argument('--data', type=str, default='data/coco.data', help='*.data file path') 278 | parser.add_argument('--weights', type=str, default='weights/yolov5s_v4.pt', help='sparse model weights') 279 | parser.add_argument('--img_size', type=int, default=416, help='inference size (pixels)') 280 | opt = parser.parse_args() 281 | print(opt) 282 | 283 | img_size = opt.img_size 284 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 285 | 286 | #the way of loading yolov5s 287 | # ckpt = torch.load(opt.weights, map_location=device) # load checkpoint 288 | # modelyolov5 = Model('models/yolov5s_v4.yaml', nc=80).to(device) 289 | # exclude = ['anchor'] # exclude keys 290 | # ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items() 291 | # if k in modelyolov5.state_dict() and not any(x in k for x in exclude) 292 | # and modelyolov5.state_dict()[k].shape == v.shape} 293 | # modelyolov5.load_state_dict(ckpt['model'], strict=False) 294 | 295 | #another way of loading yolov5s 296 | modelyolov5=torch.load(opt.weights, map_location=device)['model'].float().eval() 297 | modelyolov5.model[24].export = False # onnx export 298 | 299 | # model=modelyolov5 300 | 301 | #load yolov5s from cfg 302 | model = Darknet(opt.cfg, (img_size, img_size)).to(device) 303 | copy_weight_v4(modelyolov5,model) 304 | 305 | path='data/images/bus.jpg' 306 | img0 = cv2.imread(path) # BGR 307 | # Padded resize 308 | img = letterboxv5(img0, new_shape=416)[0] 309 | 310 | # Convert 311 | img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 312 | img = np.ascontiguousarray(img) 313 | img = torch.from_numpy(img).to(device) 314 | img = img.float() 315 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 316 | if img.ndimension() == 3: 317 | img = img.unsqueeze(0) 318 | 319 | # modelyolov5.eval() 320 | 321 | 322 | model.eval() 323 | pred = model(img)[0] 324 | 325 | pred = non_max_suppressionv5(pred, 0.4, 0.5, classes=None, 326 | agnostic=False) 327 | # Process detections 328 | for i, det in enumerate(pred): # detections per image 329 | if det is not None and len(det): 330 | # Rescale boxes from img_size to im0 size 331 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round() 332 | 333 | # Write results 334 | for *xyxy, conf, cls in det: 335 | label = '%s %.2f' % (str(int(cls)), conf) 336 | plot_one_box(xyxy, img0, label=label, color=[random.randint(0, 255) for _ in range(3)], line_thickness=3) 337 | cv2.imwrite("v5_cfg.jpg", img0) 338 | 339 | modelyolov5.eval() 340 | pred = modelyolov5(img)[0] 341 | 342 | pred = non_max_suppressionv5(pred, 0.4, 0.5, classes=None, 343 | agnostic=False) 344 | # Process detections 345 | for i, det in enumerate(pred): # detections per image 346 | if det is not None and len(det): 347 | # Rescale boxes from img_size to im0 size 348 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round() 349 | 350 | # Write results 351 | for *xyxy, conf, cls in det: 352 | label = '%s %.2f' % (str(int(cls)), conf) 353 | plot_one_box(xyxy, img0, label=label, color=[random.randint(0, 255) for _ in range(3)], 354 | line_thickness=3) 355 | cv2.imwrite("v5.jpg", img0) 356 | --------------------------------------------------------------------------------