PyTorch版YOLOv4更新了，適用於自定義數據集

機器之心報導

作者：陳萍

距離YOLO v4 的推出，已經過去 5 個多月。YOLO 框架採用 C 語言作為底層代碼，這對於慣用 Python 的研究者來說，實在是有點不友好。因此網上出現了很多基於各種深度學習框架的 YOLO 復現版本。近日，就有研究者在 GitHub 上更新了基於 PyTorch 的 YOLOv4。

從今年 4 月 YOLOv4 發布後，對於這個目標檢測框架，問的最多的問題或許就是：「有沒有同學復現 YOLOv4 的， 可以交流一下麼」。由於原版 YOLO 使用 C 語言進行編程，光憑這一點就讓不少同學望而卻步。網上有很多基於 TF/Keras 和 Caffe 等的復現版本，但不少項目只給了代碼，並沒有給出模型在 COCO、PASCAL VOC 數據集上的訓練結果。

近日，有研究者在 GitHub 上開源了一個項目：基於 PyTorch 深度學習框架的 YOLOv4 復現版本，該版本基於 YOLOv4 作者給出的實現 AlexeyAB/darknet，並在 PASCAL VOC、COCO 和自定義數據集上運行。

項目地址：https://github.com/argusswift/YOLOv4-PyTorch

除此以外，該項目還向主幹網絡添加了一些有用的注意力方法，並實現了 mobilenetv2-YOLOV4 和 mobilenetv3-YOLOV4。

attentive YOLOv4

該項目向主幹網絡添加了一些注意力方法，如 SEnet、CBAM。

SEnet (CVPR 2017)

CBAM (CVPR 2018)

mobilenet YOLOv4

該研究還實現了 mobilenetv2-YOLOV4 和 mobilenetv3-YOLOV4（只需更改 config/yolov4_config.py 中的 MODEL_TYPE 即可）。

下表展示了 mobilenetv2-YOLOV4 的性能結果：

現在我們來看該項目的詳細內容和要求。

環境要求

Nvida GeForce RTX 2080TI

CUDA10.0

CUDNN7.0

windows 或 linux 系統

python 3.6

特性

DO-Conv (https://arxiv.org/abs/2006.12030) (torch>=1.2)

Attention

fp_16 training

Mish

Custom data

Data Augment (RandomHorizontalFlip, RandomCrop, RandomAffine, Resize)

Multi-scale Training (320 to 640)

focal loss

CIOU

Label smooth

Mixup

cosine lr

安裝依賴項

運行腳本安裝依賴項。你需要提供 conda 安裝路徑（例如 ~/anaconda3）以及所創建 conda 環境的名稱（此處為 YOLOv4-PyTorch）。

需要注意的是：安裝腳本已在 Ubuntu 18.04 和 Window 10 系統上進行過測試。如果出現問題，請查看詳細的安裝說明：https://github.com/argusswift/YOLOv4-PyTorch/blob/master/INSTALL.md。

準備工作

1. git 複製 YOLOv4 庫

準備工作的第一步是複製 YOLOv4。

然後更新配置文件「config/yolov4_config.py」中「PROJECT_PATH」。

2. 數據集準備

該項目準備了 Pascal VOC 和 MSCOCO 2017 數據集。其中 PascalVOC 數據集包括 VOC 2012_trainval、VOC 2007_trainval 和 VOC2007_test，MSCOCO 2017 數據集包括 train2017_img、train2017_ann、val2017_img、val2017_ann、test2017_img、test2017_list。

PascalVOC 數據集下載命令：

MSCOCO 2017 數據集下載命令：

在數據集下載好後，需要進行以下操作：

將數據集放入目錄，更新 config/yolov4_config.py 中的 DATA_PATH 參數。

（對於 COCO 數據集）使用 coco_to_voc.py 將 COCO 數據類型轉換為 VOC 數據類型。

轉換數據格式：使用 utils/voc.py 或 utils/coco.py 將 pascal voc *.xml 格式（或 COCO *.json 格式）轉換為 *.txt 格式（Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 ...）。

3. 下載權重文件

1）darknet 預訓練權重：yolov4（https://drive.google.com/file/d/1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT/view）。

2）Mobilenet 預訓練權重：

mobilenetv2：（https://pan.baidu.com/share/init?surl=sjixK2L9L0YgQnvfDuVTJQ，提取碼：args）；

mobilenetv3：（https://pan.baidu.com/share/init?surl=75wKejULuM0ZD05b9iSftg，提取碼：args）。

3）在根目錄下創建 weight 文件夾，將下載好的權重文件放到 weight / 目錄下。

4）訓練時在 config/yolov4_config.py 中設置 MODEL_TYPE。

4. 轉換成自定義數據集（基於自定義數據集進行訓練）

1）將自定義數據集的圖片放入 JPEGImages 文件夾，將注釋文件放入 Annotations 文件夾。

2）使用 xml_to_txt.py 文件將訓練和測試文件列表寫入 ImageSets/Main/*.txt。

3）轉換數據格式：使用 utils/voc.py 或 utils/coco.py 將 pascal voc *.xml 格式（或 COCO *.json 格式）轉換為 *.txt 格式（Image_path xmin0,ymin0,xmax0,ymax0,class0 xmin1,ymin1,xmax1,ymax1,class1 ...）。

訓練

運行以下命令開始訓練，詳情參見 config / yolov4_config.py。訓練時應將 DATA_TYPE 設置為 VOC 或 COCO。

它還支持 resume 訓練，添加 --resume，使用以下命令即可自動加載 last.pt。

檢測

修改檢測圖像路徑：DATA_TEST=/path/to/your/test_data# your own images。

結果可以在 output / 中查看，如下所示：

評估（Pascal VOC 數據集）

修改評估數據集路徑：DATA_PATH=/path/to/your/test_data # your own images

評估（COCO 數據集）

修改評估數據集路徑：DATA_PATH=/path/to/your/test_data # your own images

可視化熱圖

在 val_voc.py 中設置 showatt=Ture，網絡即可輸出熱圖。

在 output / 中可以查看熱圖，如下所示：