YOLOR 訓練教學

Published in

謦伊的閱讀筆記

17 min readApr 16, 2022

YOLOR (You Only Learn One Representation) 由 ScaledYOLOv4 作者開發，能夠在與 ScaledYOLOv4 相同的準確度下，提高 88% 的速度，並且再次拿下 COCO Benchmark 排行榜冠軍。

若想了解 YOLOR 的詳細介紹可參考: YOLOR (王建堯博士介紹)

本文將要來介紹如何進行訓練，以下是 YOLOR 的 github

GitHub - WongKinYiu/yolor: implementation of paper - You Only Learn One Representation: Unified…

implementation of paper - You Only Learn One Representation: Unified Network for Multiple Tasks To get the results on…

github.com

那我們就開始吧~我的版本為 python 3.7, pytorch 1.9.0, torchvision 0.10.0

準備資料集

roboflow PlantDoc Dataset

我使用 roboflow PlantDoc Dataset: https://public.roboflow.com/object-detection/plantdoc/1，選擇 resize-416x416 YOLOv5 Pytorch 格式。

接著會得到一串指令，稍微修改一下

$ curl -L “https://public.roboflow.com/ds/yWxXqZujl9?key=4ohhBjR5iC" > plantdoc.zip; unzip -o plantdoc.zip; rm plantdoc.zip

劃分訓練/驗證集

由於 PlantDoc Dataset 的名稱較雜亂，因此將名稱重新定義，並檢查 label 內容是否為空

# 是否要重新命名 dataset 資料夾，若否則將 rename_data_name 設定為 False
rename_data_name = True
train_data_file = ‘train’if rename_data_name == True:
   org_img_path = os.path.join(train_data_file, ‘org_images’)
   img_path = os.path.join(train_data_file, ‘images’)
   org_label_path = os.path.join(train_data_file, ‘org_labels’)    
   label_path = os.path.join(train_data_file, ‘labels’)...   img_cnt = 0
   for img_file in os.listdir(org_img_path):     ...

原本的資料變為 org_images、org_labels，新命名的檔案則放在 images、labels 資料夾中

接著將資料集分為訓練/驗證集

total_data = os.listdir(os.path.join(train_data_file, “images”))
split_rate = 0.8
train_data = total_data[:math.floor(len(total_data)*split_rate)]
valid_data = total_data[math.floor(len(total_data)*split_rate):]

產生 train.txt 和 valid.txt

把剛剛分好的訓練及驗證資料集路徑寫入 train.txt 和 valid.txt 檔案

write_train_txt = os.path.join(train_data_file, ‘train.txt’)
write_valid_txt = os.path.join(train_data_file, ‘valid.txt’)train_data_list = []
with open(write_train_txt, ‘w’) as f:
    for i in train_data:
       train_data_path = os.path.join(“./images”, i) 
       f.write(‘%s\n’ % train_data_path)
       train_data_list.append(i)valid_data_list = []
with open(write_valid_txt, ‘w’) as f:
    for i in valid_data:
       valid_data_path = os.path.join(“./images”, i)
       f.write(‘%s\n’ % valid_data_path)
       valid_data_list.append(i)

產生 classes.txt

從 plantdoc/data.yaml 中可以看到總共有 30 個類別以及類別名稱

接著將類別名稱寫入 classes.txt

class_names = [‘Apple Scab Leaf’, ‘Apple leaf’, ‘Apple rust leaf’, 
                                     ...
               , ‘grape leaf black rot’, ‘grape leaf’]classes_txt = os.path.join(train_data_file, ‘classes.txt’)
class_names = [i.replace(‘ ‘, ‘_’) for i in class_names]
with open(classes_txt, ‘w’) as f:
     f.write(‘\n’.join(class_names))

產生 coco json 檔 (也可省略該步驟)

首先進行初始化

super_category = ‘leaf’dataset_info = {‘info’: {‘description’: ‘’, ‘url’: ‘’, ‘version’: ‘1.0’, ‘year’: 2022, ‘contributor’: ‘Joy’, ‘date_created’: ‘’}, ‘images’: [], ‘annotations’: [], ‘categories’: [{‘supercategory’: super_category, ‘id’: i+1, ‘name’: cls} for i, cls in enumerate(class_names)]}train_dataset = {}
valid_dataset = {}train_dataset.update(dataset_info)
valid_dataset.update(dataset_info)

將 yolo 格式轉為 coco

def yolo2coco(box_info, img_size):
    x_min = (box_info[0] — box_info[2]/2) * img_size[1]
    x_max = (box_info[0] + box_info[2]/2) * img_size[1]
    y_min = (box_info[1] — box_info[3]/2) * img_size[0] 
    y_max = (box_info[1] + box_info[3]/2) * img_size[0]    return x_min, x_max, y_min, y_max

寫入資料集的 'images' 和 'annotations' 資訊

images_dir = os.path.join(train_data_file, ‘images’)
labels_dir = os.path.join(train_data_file, ‘labels’)ann_id = 1
for idx, image_file in enumerate(os.listdir(images_dir)):
     height, width, _ = cv2.imread(os.path.join(images_dir, image_file)).shape     if image_file in train_data_list:
         dataset = train_dataset     elif image_file in valid_data_list:
         dataset = valid_dataset     dataset['images'].append({'file_name': image_file, 'width': width, 'height': height, 'id': idx+1})     ...

最後產生 annotations 的資料夾

annotation_file = os.path.join(train_data_file, ‘annotations’)if not os.path.exists(annotation_file):
    os.makedirs(annotation_file)with open(os.path.join(annotation_file, ‘train.json’), ‘w’) as json_file:
    json.dump(train_dataset, json_file)with open(os.path.join(annotation_file, ‘val.json’), ‘w’) as json_file:
    json.dump(valid_dataset, json_file)

所有資料的路徑如下

安裝環境

首先 git clone YOLOR github 並安裝需要的library

$ git clone https://github.com/WongKinYiu/yolor
$ cd yolor
$ pip install -qr requirements.txt

安裝 mish-cuda

$ git clone https://github.com/JunnYu/mish-cuda
$ cd mish-cuda
$ python setup.py build install
$ cd ..

安裝 pytorch_wavelets

$ git clone https://github.com/fbcotter/pytorch_wavelets
$ cd pytorch_wavelets
$ pip install .
$ cd ..

將剛準備好的資料集移進 yolor 資料夾中

$ mv ../plantdoc plantdoc

下載 pre-trained model

執行以下指令後會在 weights 資料夾中看到 yolor_p6.pt

$ mkdir weights
$ pip install gdown
$ gdown --id 1Tdn3yqpZ79X7R1Ql0zNlNScB1Dv9Fp76 -O "weights/yolor_p6.pt"

修改 data 中的 coco.names 和 coco.yaml

產生自定義 yaml 檔

首先查看總共有哪些類別

with open(“plantdoc/train/classes.txt”) as f:
    class_names = f.read().splitlines()
print(class_names)

將 train/ valid 檔案路徑、總類別數、類別名稱寫入 yaml 檔，其中 train/ valid 檔案路徑可寫相對路徑或絕對路徑

%%writefile data/myself_data.yaml
# train and val datasets (image directory or *.txt file with image paths)
train: /content/drive/MyDrive/yolor/plantdoc/train/train.txt
val: /content/drive/MyDrive/yolor/plantdoc/train/valid.txt
# test: ../coco/test-dev2017.txt# number of classes
nc: 30# class names
names: [‘Apple Scab Leaf’, ‘Apple leaf’, ‘Apple rust leaf’, 
                                     ...
        , ‘grape leaf black rot’, ‘grape leaf’]

產生自定義 names 檔

$ cp data/coco.names data/myself_data.names# 修改 myself_data.names 內容
names_file = ‘data/myself_data.names’
with open(names_file, ‘w’) as f:
    for i in class_names: 
        f.write(‘%s\n’ % i)

修改 cfg 中的 width, height, filters, classes

與 ScaledYOLOv4 不同的是總共有四層 yolo 層需要更改以及多加了 implicit_mul 部分

$ myself_yolor_cfg = 'cfg/myself_yolor_p6.cfg'$ cp cfg/yolor_p6.cfg {myself_yolor_cfg}# 查看原本參數
$ sed -n -e 1569p -e 1573p -e 1577p -e 1581p -e 1605p -e 1614p -e 1649p -e 1658p -e 1693p -e 1702p -e 1737p -e 1746p {myself_yolor_cfg}

這裡我將 width, height 修改為 416x416，而資料集的 label 有 30 個類別，因此 → filters=(classes + 5)*3 = (30+5)*3 = 105

$ sed -i '1569s/255/105/' {myself_yolor_cfg}$ sed -i '1573s/255/105/' {myself_yolor_cfg}$ sed -i '1577s/255/105/' {myself_yolor_cfg}$ sed -i '1581s/255/105/' {myself_yolor_cfg}$ sed -i '1605s/255/105/' {myself_yolor_cfg}$ sed -i '1614s/80/30/' {myself_yolor_cfg}$ sed -i '1649s/255/105/' {myself_yolor_cfg}$ sed -i '1658s/80/30/' {myself_yolor_cfg}$ sed -i '1693s/255/105/' {myself_yolor_cfg}$ sed -i '1702s/80/30/' {myself_yolor_cfg}$ sed -i '1737s/255/105/' {myself_yolor_cfg}$ sed -i '1746s/80/30/' {myself_yolor_cfg}# 查看修改後的參數
$ sed -n -e 1569p -e 1573p -e 1577p -e 1581p -e 1605p -e 1614p -e 1649p -e 1658p -e 1693p -e 1702p -e 1737p -e 1746p {myself_yolor_cfg}

訓練模型

準備完畢後就可以來訓練啦~~ name 是指訓練結果放置的位置，訓練完後會儲存至 runs/train/{設置的 name} 裡

$ python train.py --batch-size 8 --img 640 640 --data myself_data.yaml --cfg cfg/myself_yolor_p6.cfg --weights 'weights/yolor_p6.pt' --device 0 --name yolor_p6_0416 --hyp hyp.scratch.1280.yaml --epochs 300

若要使用多顆 GPU 進行訓練，可以執行以下指令

$ python -m torch.distributed.launch --nproc_per_node 2 --master_port 9527 train.py --batch-size 16 --img 640 640 --data data/myself_data.yaml --cfg cfg/myself_yolor_p6.cfg --weights 'weights/yolor_p6.pt' --device 0,1 --sync-bn --name yolor_p6_v1 --hyp hyp.scratch.1280.yaml --epochs 300

❗ 要訓練時遇到該問題：AttributeError: module ‘distutils’ has no attribute ‘version’

💬 可能是 setuptools 的版本過高，需要降低一下版本

import setuptools
setuptools.__version__
>> 59.8.0$ pip install setuptools == 59.5.0
# 需安裝比原本低的版本

測試模型

訓練好後，可以從 runs/train/{設置的 name} 看到 weights、訓練結果以及超參數的設置

執行以下指令來進行預測~~ 預測結果會放置在 inference/output 資料夾中

$ python detect.py --source plantdoc/test/images/00_jpg.rf.7fa2b9652948e8c39a51a68ec5c6b70a.jpg --cfg cfg/myself_yolor_p6.cfg --names data/myself_data.names --weights runs/train/yolor_p6_v1/weights/best_overall.pt --conf 0.25 --img-size 640 --device 0