這里記錄一下將yolov5用于訓練自己的數據集，這里v5的github地址

Environment

本文是直接 pip install -r requirement.txt直接安裝成功，如果不容易成功，就一個個安裝即可。如果出現報錯，多關注一下torch的版本即可。

Data preparation

label的格式:
<object-class> <x_center> <y_center> <width> <height>

舉個🌰：
414_735_0.txt

7 0.503 0.499 0.098 0.451 7 0.465 0.633 0.033 0.053

數據集格式
data
├──coco.yaml
├──coco128.yaml
├──defect_data
???├──images
??????├──train
??????├──val
???├──labels
??????├──train
??????├──val
├──…

然后制作自己的defect.yaml配置文件(直接復制coco.yaml修改或者新建一個文件將一下代碼copy進去)：

train: ./data/defect_data/images/train # absolute_path_train_img val: ./data/defect_data/images/val # absolute_path_val_img

nc: 8 # 類別數量

names: [‘quexian1’, ‘quexian2’, ... ‘quexian8’]

修改models/xxx.yaml，有s, m, l, x四個模型供選擇，需要修改類別nc數目；同時最好把anchors也修改為自己數據對應的anchors（需要自己提前聚類），不改也可train。

Pretrained Model

建議爬梯子下載最新的預訓練權重，這個repo經常更新，可能和以前的預訓練模型不適應。下載好放在weights下，在train.py里修改相應路徑。

Training

配置train.py

if __name__ == '__main__':parser = argparse.ArgumentParser()parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')parser.add_argument('--data', type=str, default='data/defect.yaml', help='data.yaml path')parser.add_argument('--epochs', type=int, default=300)parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')opt = parser.parse_args()

訓練好的模型會被保存在./runs/train/exp3/weights/last.pt和best.pt

Inference

配置detect.py

if __name__ == '__main__':parser = argparse.ArgumentParser()parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp3/weights/last.pt', help='model.pt path(s)')parser.add_argument('--source', type=str, default='data/defect_data/images/val', help='source') # file/folder, 0 for webcamparser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')opt = parser.parse_args()

如果需要inference 自己數據集原圖大小，需要修改utils/dataset.py中

class LoadImages: # for inference# Padded resizeimg = letterbox(img0, new_shape=self.img_size)[0]

修改為：

class LoadImages: # for inference# Padded resizeimg = letterbox(img0, new_shape=(math.ceil(img0.shape[0]/32)*32, math.ceil(img0.shape[1]/32)*32))[0]

Visualization of Training

利用tensorboard進行可視化，訓練開始會在主目錄生成一個runs文件，利用tensorboard打開：
tensorboard --logdir=./runs

Thanks

歡迎評論，每天回復?

總結

以上是生活随笔為你收集整理的yolov5训练自己的数据集（缺陷检测分类——看了就懂系列）的全部內容，希望文章能夠幫你解決所遇到的問題。

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