此文章只是記錄使用,以便后續查看,不作為教程,剛接觸,可能有錯誤
YOLOv5模型訓練流程
一、數據集的準備
1.在源碼根目錄新建mydata
文件夾,在此文件夾下新建images
和labels
文件夾
目錄樹如下:
├───mydata
│ ├───images
│ └───labels
2.在images下放需要標記的所有圖片,然后使用labelimg進行圖片標記
3.打開labelimg,點擊菜單欄View下的Auto Save Mode使得勾選,點擊左邊狀態欄的PascalVOC切換為YOLO,再點擊左邊狀態欄的Change Save Dir選擇路徑為上面mydata下的labels文件夾,開始標記圖片。標記完成后剪切出labels文件夾下的classes.txt到mydata目錄下
4.根目錄新建datasets/defect(隨機,但后面路徑要一致)
文件夾,在此文件夾下新建images
和labels
文件夾,還需新建的文件
夾目錄樹如下:
├───datasets
│ └───defect
│ ├───images
│ │ ├───test
│ │ ├───train
│ │ └───val
│ └───labels
│ ├───test
│ ├───train
│ └───val
5.在根目錄新建data.py
# 將圖片和標注數據按比例切分為 訓練集和測試集
import shutil
import random
import os# 原始路徑
image_original_path = "./mydata/images/"
label_original_path = "./mydata/labels/"cur_path = os.getcwd()# 訓練集路徑
train_image_path = os.path.join(cur_path, "datasets/defect/images/train/").replace(os.sep, "/")
train_label_path = os.path.join(cur_path, "datasets/defect/labels/train/").replace(os.sep, "/")# 驗證集路徑
val_image_path = os.path.join(cur_path, "datasets/defect/images/val/").replace(os.sep, "/")
val_label_path = os.path.join(cur_path, "datasets/defect/labels/val/").replace(os.sep, "/")# 測試集路徑
test_image_path = os.path.join(cur_path, "datasets/defect/images/test/").replace(os.sep, "/")
test_label_path = os.path.join(cur_path, "datasets/defect/labels/test/").replace(os.sep, "/")# 訓練集目錄
list_train = os.path.join(cur_path, "datasets/defect/train.txt").replace(os.sep, "/")
list_val = os.path.join(cur_path, "datasets/defect/val.txt").replace(os.sep, "/")
list_test = os.path.join(cur_path, "datasets/defect/test.txt").replace(os.sep, "/")train_percent = 0.6
val_percent = 0.2
test_percent = 0.2def del_file(path):for i in os.listdir(path):file_data = path + "\\" + ios.remove(file_data)def mkdir():if not os.path.exists(train_image_path):os.makedirs(train_image_path)else:del_file(train_image_path)if not os.path.exists(train_label_path):os.makedirs(train_label_path)else:del_file(train_label_path)if not os.path.exists(val_image_path):os.makedirs(val_image_path)else:del_file(val_image_path)if not os.path.exists(val_label_path):os.makedirs(val_label_path)else:del_file(val_label_path)if not os.path.exists(test_image_path):os.makedirs(test_image_path)else:del_file(test_image_path)if not os.path.exists(test_label_path):os.makedirs(test_label_path)else:del_file(test_label_path)def clearfile():if os.path.exists(list_train):os.remove(list_train)if os.path.exists(list_val):os.remove(list_val)if os.path.exists(list_test):os.remove(list_test)def main():mkdir()clearfile()file_train = open(list_train, 'w')file_val = open(list_val, 'w')file_test = open(list_test, 'w')total_txt = os.listdir(label_original_path)num_txt = len(total_txt)list_all_txt = range(num_txt)num_train = int(num_txt * train_percent)num_val = int(num_txt * val_percent)num_test = num_txt - num_train - num_valtrain = random.sample(list_all_txt, num_train)# train從list_all_txt取出num_train個元素# 所以list_all_txt列表只剩下了這些元素val_test = [i for i in list_all_txt if not i in train]# 再從val_test取出num_val個元素,val_test剩下的元素就是testval = random.sample(val_test, num_val)print("訓練集數目:{}, 驗證集數目:{}, 測試集數目:{}".format(len(train), len(val), len(val_test) - len(val)))for i in list_all_txt:name = total_txt[i][:-4]srcImage = image_original_path + name + '.png'srcLabel = label_original_path + name + ".txt"if i in train:dst_train_Image = train_image_path + name + '.png'dst_train_Label = train_label_path + name + '.txt'shutil.copyfile(srcImage, dst_train_Image)shutil.copyfile(srcLabel, dst_train_Label)file_train.write(dst_train_Image + '\n')elif i in val:dst_val_Image = val_image_path + name + '.png'dst_val_Label = val_label_path + name + '.txt'shutil.copyfile(srcImage, dst_val_Image)shutil.copyfile(srcLabel, dst_val_Label)file_val.write(dst_val_Image + '\n')else:dst_test_Image = test_image_path + name + '.png'dst_test_Label = test_label_path + name + '.txt'shutil.copyfile(srcImage, dst_test_Image)shutil.copyfile(srcLabel, dst_test_Label)file_test.write(dst_test_Image + '\n')file_train.close()file_val.close()file_test.close()if __name__ == "__main__":main()
執行代碼。
注意:根據圖片的后綴不同,修改代碼中.png,根據路徑不同修改路徑
至此,數據集準備完成。
二、配置訓練參數
1.在yolov5-master/data下新建my.yaml
# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
# COCO 2017 dataset http://cocodataset.org by Microsoft
# Example usage: python train.py --data coco.yaml
# parent
# ├── yolov5
# └── datasets
# └── coco ← downloads here (20.1 GB)# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: E:\1_Work\0802\yolov5-master\datasets\defect # dataset root dir
train: images/train # train images (relative to 'path') 118287 images
val: images/val # val images (relative to 'path') 5000 images
test: images/test # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794# Classes
nc: 3
names: ["piaochong", "qicao", "changchong"]
注意:修改對應的path,train,val,test路徑
nc為標記的類別數量
names為類別名稱,在之前的classes.txt中復制
2.在yolov5-master/models下復制yolov5s.yaml
文件重命名為yolov5s_test.yaml,修改nc類別數量
# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license# Parameters
nc: 3 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:- [10,13, 16,30, 33,23] # P3/8- [30,61, 62,45, 59,119] # P4/16- [116,90, 156,198, 373,326] # P5/32# YOLOv5 v6.0 backbone
backbone:# [from, number, module, args][[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2[-1, 1, Conv, [128, 3, 2]], # 1-P2/4[-1, 3, C3, [128]],[-1, 1, Conv, [256, 3, 2]], # 3-P3/8[-1, 6, C3, [256]],[-1, 1, Conv, [512, 3, 2]], # 5-P4/16[-1, 9, C3, [512]],[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32[-1, 3, C3, [1024]],[-1, 1, SPPF, [1024, 5]], # 9]# YOLOv5 v6.0 head
head:[[-1, 1, Conv, [512, 1, 1]],[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 6], 1, Concat, [1]], # cat backbone P4[-1, 3, C3, [512, False]], # 13[-1, 1, Conv, [256, 1, 1]],[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 4], 1, Concat, [1]], # cat backbone P3[-1, 3, C3, [256, False]], # 17 (P3/8-small)[-1, 1, Conv, [256, 3, 2]],[[-1, 14], 1, Concat, [1]], # cat head P4[-1, 3, C3, [512, False]], # 20 (P4/16-medium)[-1, 1, Conv, [512, 3, 2]],[[-1, 10], 1, Concat, [1]], # cat head P5[-1, 3, C3, [1024, False]], # 23 (P5/32-large)[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)]
3.修改train.py:
修改項 | 值 | 解釋 |
---|---|---|
--weights | default=ROOT / ‘yolov5s.pt’ | 預訓練權重,在官方Github下載 |
--cfg | default=‘models/yolov5s_test.yaml’ | 加載模型,之前復制修改的文件 |
--data | default=ROOT / ‘data/my.yaml’ | 配置文件 |
--epochs | default=2 | 訓練批次 |
--batch-size | default=2 | 每批次的輸入數據量 |
修改完成之后運行即可開始訓練。
三、檢驗訓練的模型的推理效果
修改detect.py文件:
修改項 | 值 | 解釋 |
---|---|---|
--weights | default=ROOT / ‘best.pt’ | 為訓練完成后在runs/train/exp/weights下的best.pt文件 |
--source | default=ROOT / ‘datasets/defect/images/test’ | 測試數據目錄 |
執行文件即可
結果在yolov5-master/runs/detect/exp下。
注意:如果生成的圖片沒有框,可以降低置信度排查錯誤,即detect.py文件中的–conf-thres,–iou-thres。