本文參考github上大神的開源剪枝項目進行學習與分享，具體鏈接放在文后，希望與大家多多交流！

一、原模型訓練

在官方源碼上訓練yolov5模型，支持v6.0分支的n/s/m/l模型，我這里使用的是v5s，得到后將項目clone到本機上

git clone https://github.com/midasklr/yolov5prune.git

cd進入文件夾后，新建runs文件夾，將訓練好的模型放入runs/your_train/weights/xxx.pt，我的原模型map0.5:0.95為0.84左右，模型與data.yaml設置好后可以進行稀疏化訓練了。

二、稀疏化訓練

python train_sparity.py --st --sr 0.0002 --weights yolov5s.pt --data data/your_data.yaml --epochs 100 --imgsz 512 --adam ...

注意：1、若原模型訓練時未使用adam，則這里也不要使用adam。2:data.yaml文件改成自己的數據集文件。3：這里sr參數為稀疏化系數，具體值為多少根據不同的數據集和模型，一般設置不一樣，需要自己多試試，比如我的是單類別目標檢測，設置為0.0002時幾乎不變化（如下圖）。

這里我們cd到runs文件夾路徑，然后輸入tensorboard實時監看訓練過程，logdir也指向runs文件夾即可，然后打開網頁輸入你的ip，端口號一般為6006，比如192.168.xx.xx:6006就可以監看。

tensorboard --logdir=/home/user2/yinjiacheng/pytorch-ssd/pytorch-ssd/runs --host=0.0.0.0

這個直方圖的縱軸代表訓練次數（從上往下訓練次數增加），橫軸的峰值應隨著訓練次數不斷逼近0軸，代表著大多數bn已經變得稀疏，而出現下圖這個情況或者稀疏過慢(逼近0軸的過程緩慢)時，代表sr值應該適量增大。

所以我調整了sr為0.02，但此時bn收斂過快，且mAP下降嚴重，如下圖所示，表示sr值調的過大了，需要適當減少，直至直方圖逼近0軸的同時，mAP與原模型相差不大。

然后我將sr值調整為0.01，此時的bn直方圖與mAP0.5曲線如下圖所示：

可以看出bn直方圖往0軸逼近的速度不是太快也不是太慢(好像在說廢話)，而mAP雖然有掉點但還是可以接受，最后可以finetune回來一些，所以可以選擇這一稀疏化模型進行下一步的剪枝。

三、剪枝

python prune.py --percent 0.5 --weights runs/train/exp/weights/last.pt --data data/your_data.yaml --cfg models/yolov5s.yaml

這里percent參數是剪枝比例，weights選擇剛剛訓練好的稀疏化模型，cfg選擇你使用的yaml即可，剪枝完成后主文件夾下會出現pruned_model.pt文件，這就是剪枝后的模型，但會發現怎么模型大小下降的不多，因為它還是FP32精度的，而我們yolov5訓練完成后的模型為FP16精度，所以其實大小還是下降挺多的（最后我微調后得到的是5.8M,而原模型為13.6M）。

四、微調

python finetune_pruned.py --weights pruned_model.pt --data data/your_data.yaml --epochs 150 --imgsz 320 --adam

這里還是一樣，若前面沒有使用adam則這里也不要用，imgsz調成自己的輸出尺寸，200個epoch還是finetune回來不少的。

最后我獲得的模型為5.48M，參數量Param為2.74M，輸入尺寸為320*320時FLOPs為2.17G，比原模型大小13.6M，FLOPs為3.9G明顯減小。

五、detect

該項目中的detect.py和detectpruned.py代碼有些問題，大家可以把微調后的pt模型放到官方項目中去detect，但有幾點需要改。

現在我們是兩個項目，一個是剪枝的yolov5_prune，一個是官方的yolov5，將yolov5_prune/model/pruned_common.py文件復制到官方yolov5/model/ 中去。

將pruned_common.py中第26行的

from utils.plots import Annotator, colors, save_one_box

save_one_box刪掉，再在25行最后加上，因為在官方文件中save_one_box這個def是在utils下的general中定義的。

from utils.general import (LOGGER, check_requirements, check_suffix, check_version, colorstr, increment_path,make_divisible, non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh,save_one_box)

打開官方yolov5/utils/general.py，在第43行后，也就是聲明第一個class前加上以下代碼

def set_logging(name=None, verbose=True):# Sets level and returns loggerfor h in logging.root.handlers:logging.root.removeHandler(h) # remove all handlers associated with the root logger objectrank = int(os.getenv('RANK', -1)) # rank in world for Multi-GPU trainingslogging.basicConfig(format="%(message)s", level=logging.INFO if (verbose and rank in (-1, 0)) else logging.WARNING)return logging.getLogger(name)LOGGER = set_logging(__name__) # define globally (used in train.py, val.py, detect.py, etc.)

最后根據設置運行detect.py即可，我在detect中加了將漏報和誤報圖片分別放在不同文件夾的功能，最后在5000張val集中分別多了幾十張吧，還算可以接受。

python detect.py --weights runs/train/yolo5s_face_prune/weights/last.pt --conf 0.5 --img-size 320 --source data/your_data/images/val --save-txt

參考鏈接：midasklr/yolov5prune at v6.0 (github.com)

總結

以上是生活随笔為你收集整理的yolov5s模型剪枝详细过程(v6.0)的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。