Ubuntu使用官方Yolov5训练自己的数据集（小白向）

一、下载官方代码

本文使用的代码是2021年2月份的。最新代码链接：
github(https://github.com/ultralytics/yolov5)

二、数据集的制作

数据集的文件分级如下所示
文件结构必须是这样的，images放图，labels放标签

|————shujuji
	|————train
		|————images
		|————labels
	|————val
		|————images
		|————labels

yolov5使用的label是txt格式的，可以使用标注工具LabelImg（GitHub搜索即可）

三、文件修改

1、models文件夹

models文件夹中的yaml文件记录了网络的配置信息，分为yolov5s、yolov5m、yolov5l、yolov5x四个，s是最轻量的版本网络深度最小，其他的依次增大，以yolov5s.yaml为例：

# parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple

# anchors
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 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 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, 9, 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, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, C3, [1024, False]],  # 9
  ]

# YOLOv5 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)
  ]

nc是要类别数量，根据自己的数据集中的分类数修改nc即可，其他的无需改动。

2、data文件夹

data文件夹中的yaml文件记录着数据集的地址和类别数及类别名，以coco128.yaml为例

# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
#   /parent_folder
#     /coco128
#     /yolov5

 这步是从coco上下载数据集，训练自己数据集时注释掉
# download command/URL (optional)
# download: 
# https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/  # 128 images
val: ../coco128/images/train2017/  # 128 images

本文的路径可改为
train: ../shujuji/train/images
val: ../shujuji/val/images

# number of classes
nc: 80

# class names
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
         'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
         'hair drier', 'toothbrush' ]

train和val是训练集和验证集图片的路径，可以是绝对路径也可以是相对路径(yolov5可以通过图片路径推断出label路径)。nc为类别数。names为类别名称，根据自己的数据集进行修改。

至此，所需修改的文件就完成了。

四、命令输入

在Ubuntu环境下，shell中输入

python train.py --help

可以查看可以传入的参数类型

其中weights是权重文件 .pt 格式，可以输入空格，代表使用随机权重，或者输入权重文件的路径
cfg是模型的yaml文件，一般存放在models文件夹里
data是数据集的yaml文件，一般存放在data文件夹里
epochs是训练轮数，默认300轮
batch-size是batch数，默认16
img是输入图片大小，网络会自动按参数进行resize,默认640X640

示例：

cd yolov5-master
python train.py --img 640 --batch 16 --epochs 300 --data ./data/data.yaml --cfg ./models/yolov5s.yaml --weights '' 

这样就大功告成了