目录
项目简介
基于PyTorch实现Yolov5算法，作为Yolov5算法的复现。可以帮助读者更好的理解它的网络结构、训练流程、损失计算等。读者也可以使用该仓库训练自己的数据集，项目代码可在github获取。
Yolov5算法介绍，参考博客：零基础Yolov5学习-CSDN博客
项目github地址：GitHub - wzl639/yolov5-pytorch
数据集加载
数据集加载类定义在dataset.yolo_dataset.py中。__getitem__()方法主要包括：1）对图片进行数据增强操作，包括：mosaic数据增强、缩放、色域变换、旋转；2）预处理：维度转换、归一化；3）标签处理：voc原始框坐标是坐上右下形式，转换为中心点宽高形式[cx, cy, w, h]，并且是归一化的形式
class YoloDataset(Dataset):
def __init__(self, annotation_lines, input_shape, num_classes, epoch_length, mosaic, train, mosaic_ratio=0.7):
super(YoloDataset, self).__init__()
self.annotation_lines = annotation_lines # 图片标注信息，.txt文件
self.input_shape = input_shape # 输入模型图片大小
self.num_classes = num_classes #
self.epoch_length = epoch_length
self.mosaic = mosaic # 是否使用mosaic数据增强
self.train = train # 训练数据还是测试数据
self.mosaic_ratio = mosaic_ratio # 使用mosaic数据增强的比例，前多少个epoch使用
self.epoch_now = -1 # 当前训练批次，用来控制是否使用mosaic数据增强， 随训练过程修改
self.length = len(self.annotation_lines) # 数据集大小
def __len__(self):
return self.length
def __getitem__(self, index):
"""
获取单个数据和标签, 训练时进行数据的随机增强，验证时不进行数据的随机增强
数据增强主要包括：mosaic数据增强、缩放、色域变换、旋转，预处理：维度转换、归一化
标签处理：voc原始框坐标是坐上右下形式，转换为中心点宽高形式[cx, cy, w, h]，并且是归一化的形式
index: 获取数据的索引
return:
image：处理后图片，np格式
box：图片对应标签，np格式
"""
index = index % self.length
if self.mosaic:
if self.rand() < 0.5 and self.epoch_now < self.epoch_length * self.mosaic_ratio:
lines = sample(self.annotation_lines, 3)
lines.append(self.annotation_lines[index])
shuffle(lines)
image, box = self.get_random_data_with_Mosaic(lines, self.input_shape)
else:
image, box = self.get_random_data(self.annotation_lines[index], self.input_shape, random=self.train)
else:
image, box = self.get_random_data(self.annotation_lines[index], self.input_shape, random=self.train)
image = np.transpose(preprocess_input(np.array(image, dtype=np.float32)), (2, 0, 1))
box = np.array(box, dtype=np.float32)
if len(box) != 0:
box[:, [0, 2]] = box[:, [0, 2]] / self.input_shape[1]
box[:, [1, 3]] = box[:, [1, 3]] / self.input_shape[0]
box[:, 2:4] = box[:, 2:4] - box[:, 0:2]
box[:, 0:2] = box[:, 0:2] + box[:, 2:4] / 2
return image, box模型定义
模型相关的定义在models在模块中，CSPdarknet.py中定义了主干网络、yolo.py是整个模型的定义。
class YoloBody(nn.Module):
def __init__(self, anchors_mask, num_classes, phi):
super(YoloBody, self).__init__()
depth_dict = {'s' : 0.33, 'm' : 0.67, 'l' : 1.00, 'x' : 1.33,}
width_dict = {'s' : 0.50, 'm' : 0.75, 'l' : 1.00, 'x' : 1.25,}
dep_mul, wid_mul = depth_dict[phi], width_dict[phi]
base_channels = int(wid_mul * 64) # 64
base_depth = max(round(dep_mul * 3), 1) # 3
# 输入图片是640, 640, 3，初始的基本通道是64
# 生成CSPdarknet53的主干模型
# 获得三个有效特征层，他们的shape分别是：
# 80,80,256
# 40,40,512
# 20,20,1024
#---------------------------------------------------#
self.backbone = CSPDarknet(base_channels, base_depth)
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
self.conv_for_feat3 = Conv(base_channels * 16, base_channels * 8, 1, 1)
self.conv3_for_upsample1 = C3(base_channels * 16, base_channels * 8, base_depth, shortcut=False)
self.conv_for_feat2 = Conv(base_channels * 8, base_channels * 4, 1, 1)
self.conv3_for_upsample2 = C3(base_channels * 8, base_channels * 4, base_depth, shortcut=False)
self.down_sample1 = Conv(base_channels * 4, base_channels * 4, 3, 2)
self.conv3_for_downsample1 = C3(base_channels * 8, base_channels * 8, base_depth, shortcut=False)
self.down_sample2 = Conv(base_channels * 8, base_channels * 8, 3, 2)
self.conv3_for_downsample2 = C3(base_channels * 16, base_channels * 16, base_depth, shortcut=False)
self.yolo_head_P3 = nn.Conv2d(base_channels * 4, len(anchors_mask[2]) * (5 + num_classes), 1)
self.yolo_head_P4 = nn.Conv2d(base_channels * 8, len(anchors_mask[1]) * (5 + num_classes), 1)
self.yolo_head_P5 = nn.Conv2d(base_channels * 16, len(anchors_mask[0]) * (5 + num_classes), 1)
def forward(self, x):
# backbone
feat1, feat2, feat3 = self.backbone(x)
P5 = self.conv_for_feat3(feat3)
P5_upsample = self.upsample(P5)
P4 = torch.cat([P5_upsample, feat2], 1)
P4 = self.conv3_for_upsample1(P4)
P4 = self.conv_for_feat2(P4)
P4_upsample = self.upsample(P4)
P3 = torch.cat([P4_upsample, feat1], 1)
P3 = self.conv3_for_upsample2(P3)
P3_downsample = self.down_sample1(P3)
P4 = torch.cat([P3_downsample, P4], 1)
P4 = self.conv3_for_downsample1(P4)
P4_downsample = self.down_sample2(P4)
P5 = torch.cat([P4_downsample, P5], 1)
P5 = self.conv3_for_downsample2(P5)
# 第三个特征层
# y3=(batch_size,75,80,80)
out2 = self.yolo_head_P3(P3)
# 第二个特征层
out1 = self.yolo_head_P4(P4)
# 第一个特征层
out0 = self.yolo_head_P5(P5)
return out0, out1, out2损失计算
损失函数定义model.yolo_loss.py中，主要流程：首先对网络输出结果进行shape调整、sigmoid，方便后续损失计算，然后调用get_target()来获取网络应该得到的输出结果y_true（这里包含正样本匹配的过程），最后网络输出和y_true求损失。其中比较难的是get_target()函数。
class YOLOLoss(nn.Module):
def __init__(self, anchors, num_classes, input_shape, cuda, anchors_mask=[[6, 7, 8], [3, 4, 5], [0, 1, 2]],
label_smoothing=0):
super(YOLOLoss, self).__init__()
def forward(self, l, input, targets=None):
"""
yolov5单层输出损失计算
l: 代表使用的是第几个有效特征层
input: 模型当前层输出， bs, 3*(5+num_classes), 13, 13
targets: 真实框的标签情况 list [batch_size, num_gt, 5]
return: 当前层计算得到的损失
"""
# 获得当前批次图片数量，特征层的高和宽
bs = input.size(0)
in_h = input.size(2)
in_w = input.size(3)
# 计算步长， stride_h = stride_w = 32、16、8
stride_h = self.input_shape[0] / in_h
stride_w = self.input_shape[1] / in_w
# 将原图anchor缩放到特征图大小，此时获得的scaled_anchors大小是相对于特征层的
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors]
# 调整模型输出，将每个预测框的预测信息拆分出来
prediction = input.view(
bs, len(self.anchors_mask[l]),
self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous() # [b, 3, 20, 20, 25(5 + num_classes)]
# 先验框的中心位置的调整参数
x = torch.sigmoid(prediction[..., 0]) # [b, 3, 20, 20]
y = torch.sigmoid(prediction[..., 1])
# 先验框的宽高调整参数
w = torch.sigmoid(prediction[..., 2])
h = torch.sigmoid(prediction[..., 3])
# 获得置信度，是否有物体
conf = torch.sigmoid(prediction[..., 4])
# 种类置信度
pred_cls = torch.sigmoid(prediction[..., 5:])
# 获得网络应该有的预测结果, 这里包含正样本匹配的过程
# 网络应该有的预测结果:y_true: [1, 3, 20, 20, 25(5 + num_classes)]，用于后续loss计算
# noobj_mask: [1, 3, 20, 20] noobj_mask代表无目标的特征点，暂时没有用到
y_true, noobj_mask = self.get_target(l, targets, scaled_anchors, in_h, in_w)
# 将预测结果进行解码, 方便后面计算giou损失
pred_boxes = self.get_pred_boxes(l, x, y, h, w, targets, scaled_anchors, in_h, in_w) # [1, 3, 20, 20, 4]
if self.cuda:
y_true = y_true.cuda()
# noobj_mask = noobj_mask.cuda()
# loss计算
loss = 0
n = torch.sum(y_true[..., 4] == 1) # 统计当前批次数据中是否有正样本
if n != 0:
# 当前batch数据有目前，计算正样本的位置和类别损失
giou = self.box_giou(pred_boxes, y_true[..., :4]) # [1, 3, 20, 20]
loss_loc = torch.mean((1 - giou)[y_true[..., 4] == 1]) # 只用正样本anchor计算
loss_cls = torch.mean(self.BCELoss(pred_cls[y_true[..., 4] == 1],
self.smooth_labels(y_true[..., 5:][y_true[..., 4] == 1],
self.label_smoothing,
self.num_classes))) # 只用正样本anchor计算
loss += loss_loc * self.box_ratio + loss_cls * self.cls_ratio
# 计算置信度的标签，这里就是将正样本anchor预测框和真实框的giou值作为置信度，giou值越大执行度越大
# torch.where(condition, x, y), 若满足条件，则取x中元素 若不满足条件，则取y中元素
tobj = torch.where(y_true[..., 4] == 1, giou.detach().clamp(0), torch.zeros_like(y_true[..., 4]))
else:
# 当前batch数据没有目标，只计算置信度损失
tobj = torch.zeros_like(y_true[..., 4]) # [1, 3, 20, 20] tobj是物体置信度label
loss_conf = torch.mean(self.BCELoss(conf, tobj))
loss += loss_conf * self.balance[l] * self.obj_ratio
# if n != 0:
# print(loss_loc * self.box_ratio, loss_cls * self.cls_ratio, loss_conf * self.balance[l] * self.obj_ratio)
return loss模型训练
数据集准备
项目使用VOC数据集作为例子
VOC数据集介绍
PASCAL VOC 挑战赛是一个世界级的计算机视觉挑战，主要包括以下几类：图像分类(Object Classification)，目标检测(Object Detection)，目标分割(Object Segmentation)，行为识别(Action Classification) 。数据集中包括了20个常见的目标类别，例如人、汽车、猫、狗等，包含如下几个目录：
Annotations：这个文件夹内主要存放了数据的标签，里面包含了每张图片的bounding box信息，主要用于目标检测。
ImageSets：用于存放不同任务的划分的数据集。
Segmentation：只包含一组 [train.txt, trainval.txt, val.txt] 文件，各文件只有1列，为图片名称。
JPEGImages：这里存放的就是JPG格式的原图，包含17125张彩色图片，但只有一部分(2913张)是用于分割的。
SegmentationClass：语义分割任务中用到的label图片，PNG格式，共2913张，与原图的每一张图片相对应。
SegmentationObject：实例分割任务用到的label图片，在语义分割中用不到。
VOC数据集官网下载地址：The PASCAL Visual Object Classes Homepage
VOC数据集包含目标检测和分割标注，本仓库只需要用到目标检测部分，需要用到下面几个文件夹数据：Annotations、ImageSets和JPEGImages。
预处理
数据集下载好需要用voc_annotation.py脚本来划分训练验证和标签预处理，修改脚本中VOCdevkit_path指向数据集目录，该脚本执行完会，会在数据集目录创建2007_train.txt和2007_val.txt，用于训练。也可以直接从百度网盘下载作者处理好的数据集：
百度网盘地址: https://pan.baidu.com/s/1MF5e8wgdkJ6kFjnNhhLfXA?pwd=dtcr
提取码: dtcr
训练
train.py参数详解：
if __name__ == '__main__':
# argparse模块，当字典一样用，方便传参
parser = argparse.ArgumentParser(description="----------------yolov5 train-----------------")
parser.add_argument('--cuda', default='True', help='use cuda')
# 模型损失相关
parser.add_argument('--classes_path', default='./model_data/voc_classes.txt', help='location of classes path')
parser.add_argument('--anchors_path', default='./model_data/yolo_anchors.txt', help='location of anchors path')
parser.add_argument('--anchors_mask', default="[[6, 7, 8], [3, 4, 5], [0, 1, 2]]", help='')
parser.add_argument('--model_path', default='./model_data/yolov5_s.pth', help='location of model path')
parser.add_argument('--phi', default="s", help='model type s, m, l, or x')
parser.add_argument('--label_smoothing', default=0.0, type=float, help='label smoothing value')
# 数据集相关
parser.add_argument('--num_workers', default=4, type=int, help='nums of data load thread')
parser.add_argument('--mosaic', default='True', help='use mosaic data enhancement or not')
parser.add_argument('--input_shape', default=640, type=int, help='model input size')
parser.add_argument('--train_annotation_path', default='./data/VOC2007/2007_train.txt', help='location of train annotation path')
parser.add_argument('--val_annotation_path', default="./data/VOC2007/2007_val.txt", help='location of val annotation path')
# 优化器相关
parser.add_argument('--Init_lr', default=1e-2, type=float, help='train total epochs')
parser.add_argument('--Min_lr', default=0.00001, type=float, help='batch size')
parser.add_argument('--optimizer_type', default="sgd", help='optimizer type, adam or sgd')
parser.add_argument('--momentum', default=0.937, type=float, help='momentum for optimizer')
parser.add_argument('--weight_decay', default=5e-4, help='weight decay')
parser.add_argument('--lr_decay_type', default="cos", help='lr decay , cos or step')
# 训练相关
parser.add_argument('--epochs', default=100, type=int, help='train total epochs')
parser.add_argument('--batch_size', default=4, type=int, help='batch size')
parser.add_argument('--save_period', default=1, type=int, help='save period')
parser.add_argument('--save_dir', default='./logs/', help='location of checkpoint')
args = parser.parse_args()
# 参数转换
args.anchors_mask = eval(args.anchors_mask)
args.mosaic = eval(args.mosaic)
args.cuda = eval(args.cuda)
args.input_shape = [args.input_shape, args.input_shape]
# print(args, type(args.cuda))
# 调用主函数
main(args)说明:
mosaic数据增强：参考YoloX，由于Mosaic生成的训练图片，远远脱离自然图片的真实分布。本代码会在训练结束前的N个epoch自动关掉Mosaic100个世代会关闭30个世代（比例可在dataloader.py调整）
label_smoothing 标签平滑：一般0.01以下。如0.01、0.005
optimizer_type ：优化器种类可选的有adam、sgd，当使用Adam优化器时建议设置 Init_lr=1e-3，当使用SGD优化器时建议设置 Init_lr=1e-2，adam会导致weight_decay错误，使用adam时建议设置为0。
模型推理
predict.py支持单张图片预测、视频检测、和目录遍历检测等功能,通过指定mode进行模式的修改。参数如下：
# argparse模块
parser = argparse.ArgumentParser(description="----------------yolov5 predict.py-----------------")
# 模型相关
parser.add_argument('--model_path', default='./model_data/yolov5_s.pth', help='location of model path')
parser.add_argument('--classes_path', default='./model_data/coco_classes.txt', help='location of classes path')
parser.add_argument('--anchors_path', default='./model_data/yolo_anchors.txt', help='location of anchors path')
parser.add_argument('--anchors_mask', default="[[6, 7, 8], [3, 4, 5], [0, 1, 2]]", help='')
parser.add_argument('--input_shape', default=640, type=int, help='model input size')
parser.add_argument('--phi', default="s", help='model type s, m, l, or x')
parser.add_argument('--confidence', default=0.5, type=float, help='object confidence')
parser.add_argument('--nms_iou', default=0.3, type=float, help='nms iou threshold')
parser.add_argument('--letterbox_image', default='True', help='use letterbox for input image or not')
parser.add_argument('--cuda', default='True', help='use cuda')
# predict表示单张图片预测,video表示视频检测, dir_predict表示遍历文件夹进行检测并保存
parser.add_argument('--mode', default='predict', help='input mode, predict, video, or dir_predict')
# crop指定了是否在单张图片预测后对目标进行截取, crop仅在mode='predict'时有效
parser.add_argument('--img_path', default='./data/VOC2007/JPEGImages/2007_000027.jpg', help='location of image path')
parser.add_argument('--crop', default='False', help='')
# video_path用于指定视频的路径,video_save_path表示视频保存的路径，当video_save_path=""时表示不保存
# video_fps用于保存的视频的fps,video_path、video_save_path和video_fps仅在mode='video'时有效
parser.add_argument('--video_path', default="test.mp4", help='')
parser.add_argument('--video_save_path', default='test_result.mp4', help='')
parser.add_argument('--video_fps', default=25, type=int, help='')
# dir_origin_path指定了用于检测的图片的文件夹路径,dir_save_path指定了检测完图片的保存路径
# dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
parser.add_argument('--dir_origin_path', default="./imgs/", help='')
parser.add_argument('--dir_save_path', default='./result/', help='')
args = parser.parse_args()
# 参数转换
args.anchors_mask = eval(args.anchors_mask)
args.input_shape = [args.input_shape, args.input_shape]
args.letterbox_image = eval(args.letterbox_image)
args.cuda = eval(args.cuda)
args.crop = eval(args.crop)
print(args)
# 调用主函数
main(args)运行单张检测结果：
模型测试
目标检测问题，一般的常用评价指标有：
精度评价指标：mAP（mean Average Precision，平均准确度均值），平均正确率（AP），准确率 （Accuracy），精确率（Precision），召回率（Recall）。
速度评价指标：FPS（即每秒处理的图片数量或者处理每张图片所需的时间，在同一硬件条件下进行比较）
eval.py脚本可以测试模型再VOC数据上的指标。
# argparse模块
parser = argparse.ArgumentParser(description="----------------yolov5 predict.py-----------------")
# 模型相关
parser.add_argument('--model_path', default='./model_data/yolov5_s.pth', help='location of model path')
parser.add_argument('--classes_path', default='./model_data/coco_classes.txt', help='location of classes path')
parser.add_argument('--anchors_path', default='./model_data/yolo_anchors.txt', help='location of anchors path')
parser.add_argument('--anchors_mask', default="[[6, 7, 8], [3, 4, 5], [0, 1, 2]]", help='')
parser.add_argument('--input_shape', default=640, type=int, help='model input size')
parser.add_argument('--phi', default="s", help='model type s, m, l, or x')
parser.add_argument('--confidence', default=0.5, type=float, help='object confidence')
parser.add_argument('--nms_iou', default=0.3, type=float, help='nms iou threshold')
parser.add_argument('--letterbox_image', default='True', help='use letterbox for input image or not')
parser.add_argument('--cuda', default='True', help='use cuda')
# MINOVERLAP用于指定想要获得的mAP0.x,比如计算mAP0.75，可以设定MINOVERLAP = 0.75。
parser.add_argument('--MINOVERLAP', default=0.5, type=float, help='')
# 指向VOC数据集所在的文件夹, 默认指向根目录下的VOC数据集
parser.add_argument('--VOCdevkit_path', default='./data/VOC2007/', help='')
# map_vis用于指定是否开启VOC_map计算的可视化
parser.add_argument('--map_vis', default='False', help='')
# 结果输出的文件夹，默认为map_out
parser.add_argument('--map_out_path', default='map_out', help='')
args = parser.parse_args()
# 参数转换
args.anchors_mask = eval(args.anchors_mask)
args.input_shape = [args.input_shape, args.input_shape]
args.letterbox_image = eval(args.letterbox_image)
args.cuda = eval(args.cuda)
args.map_vis = eval(args.map_vis)
print(args)
# 调用主函数
main(args)测试结果会保存在--map_out_path指定的文件夹下，下面是每个类的AP和所有类别的mAP:
扫一扫
864
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
