YOLOv5代码阅读笔记

本来想先总结yolov5的各种知识点，但是看了一位大佬发的博客，瞬间就跪了，链接放上：
深入浅出Yolo系列之Yolov5核心基础知识完整讲解代码看完一遍后，感觉理解还不够深刻，决定近期再把代码过一遍，顺便写个阅读笔记加深记忆。
看代码建议从推理部分开始看。

一、detect.py

由于以前就是用的v5团队写的pytorch版yolov3，detect.py跟v3的代码基本一样，还是原来的配方。这部分的代码很简单，认真读一下基本就懂了，说一下其中的几个知识点：

１、cudnn.benchmark = True

在推理时使用torch.backends.cudnn.benchmark = true，可以让内置的 cuDNN 的 auto-tuner 自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题。
一般来讲，应该遵循以下准则：
①. 如果网络的输入数据维度或类型上变化不大，设置 torch.backends.cudnn.benchmark = true 可以增加运行效率；
②. 如果网络的输入数据在每次 iteration 都变化的话，会导致 cnDNN 每次都会去寻找一遍最优配置，这样反而会降低运行效率。
所以，推理视频流时，应保证每个摄像头传递给算法的图片一样大小。

２、pred = model(img, augment=opt.augment)[0]

pred的shape是(1, num_boxes, 5+num_class)
num_boxes为模型在３个特征图上预测出来的框的个数，num_boxes = h/32 * w/32 + h/16 * w/16 + h/8 * w/8
5+nnum_class的值为：
pred[…, 0:4]为预测框坐标,预测框坐标为xywh(中心点+宽长)格式
pred[…, 4]为objectness置信度
pred[…, 5:-1]为分类结果
预测得到的这一堆框送入后面的NMS函数。

３、def non_max_suppression()

# Detections matrix nx6 (xyxy, conf, cls)
if multi_label:
    # i:从0开始的序号, j:模型预测的框的conf>阀值的类别
    i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
    # box[i]第i个框的坐标，x[i, j + 5, None]:置信度，j[:, None]:类别
    x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)

else:  # best class only
    conf, j = x[:, 5:].max(1, keepdim=True)
    x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]

当检测种类>1时，使用multi_label，作者的注释表明会增加0.5ms/img，个人没感觉到这两种代码的计算量有啥明显的区别，都要做>conf_thres这一步。
.T转置
.t() 是 .transpose函数的简写版本，但只能对2维以下的tensor进行转置，
.transpose函数对一个n维tensor交换其任意两个维度的顺序，
而 .T 是 .permute 函数的简化版本，不仅可以操作2维tensor，甚至可以对n维tensor进行转置。当然当维数n=2时，.t() 与 .T 效果是一样的。

i = torch.ops.torchvision.nms(boxes, scores, iou_thres)  

boxes (Tensor[N, 4])) – bounding boxes坐标. 格式：(x1, y1, x2, y2)
scores (Tensor[N]) – bounding boxes得分
iou_threshold (float) – IoU过滤阈值
返回经iou_thres过滤scores后的boxes索引

# 加权NMS
if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
    try:  # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
        iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
        weights = iou * scores[None]  # box weights
        x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
        if redundant:
            i = i[iou.sum(1) > 1]  # require redundancy

这里差点就当成了soft nms，加权非极大值抑制与传统的非极大值抑制相比，是在进行矩形框剔除的过程中，并未将那些与当前矩形框IOU大于阈值，且类别相同的框直接剔除，而是根据网络预测的置信度进行加权，得到新的矩形框，把该矩形框作为最终预测的矩形框，再将那些框剔除。

４、strip_optimizer(opt.weights)

该函数主要是清除掉权重内的opt等信息。

二、class LoadImages

该类的主要作用是推理时读取数据，作者已整合了图片、视频、文件夹、摄像头的读取，非常方便。

１、img = letterbox(img0, new_shape=self.img_size)[0]

对图片进行resize+pad，letterbox的默认参数里：auto=True, scaleFill=False, scaleup=True

# Scale ratio (new / old)   # 计算缩放因子
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
if not scaleup:  # only scale down, do not scale up (for better test mAP)
    r = min(r, 1.0)

缩放到输入大小img_size的时候，如果没有设置上采样的话，则只进行下采样，因为上采样图片会让图片模糊，对训练不友好影响性能。

# Compute padding
ratio = r, r  # width, height ratios
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
if auto:  # minimum rectangle   # 获取最小的矩形填充
    dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding
elif scaleFill:  # stretch  
    dw, dh = 0.0, 0.0
    new_unpad = (new_shape[1], new_shape[0])
    ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

如果scaleFill=True,则不进行填充，直接resize成img_size,任由图片进行拉伸和压缩。

auto=True表示推理时不是将图片填充为正方形，而是获取最小的矩形填充。
此处的计算流程为：先按以前的正方形填充方式，以图片长边为基准，计算宽高resize到设置的尺寸640的缩放因子，原图按该缩放因子进行缩放，长边resize为640，短边比640小，再用640－缩放后长度，即为宽高的padding，dw, dh = np.mod(dw, 64), np.mod(dh, 64) 因为是在两侧加padding，且要求图片必须为32的倍数，故用64对padding取余，得到的dw,dh为宽高最终的padding。

top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))

由于padding可能为带有0.5的小数，故加减0.1用rond()函数四舍五入取值后再取整。

２、img = np.ascontiguousarray(img)

将数组内存转为连续，提高运行速度，不转的话也可能会报错

三、train.py

1 、main()

opt参数解析：
weights:加载预训练的yolov5权重文件
cfg:模型配置文件，网络结构
data:数据集配置文件，数据集路径，类名等
hyp:超参数文件
epochs:训练总轮次
batch-size:批次大小
img-size:输入图片分辨率大小
rect:是否采用矩形训练，默认False
resume:接着打断训练上次的结果接着训练
nosave:只保存最后一个模型，默认False
notest:只对最后一个epoch进行test，默认False
noautoanchor:不自动调整anchor，默认False
evolve:是否进行超参数进化，默认False
bucket:谷歌云盘bucket，一般不会用到
cache-images:是否提前缓存图片到内存，以加快训练速度，默认False
image-weights:使用加权图像选择进行训练???
name:数据集名字，如果设置：results.txt to results_name.txt，默认无
device:训练的设备，cpu；0(表示一个gpu设备cuda:0)；0,1,2,3(多个gpu设备)
multi-scale:是否进行多尺度训练，默认False
single-cls:数据集是否只有一个类别，默认False
adam:是否使用adam优化器
sync-bn:是否使用跨卡同步BN,在DDP模式使用
local_rank:gpu编号
logdir:存放日志的目录
workers:dataloader的最大worker数量

# Set DDP variables
opt.total_batch_size = opt.batch_size
opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1    # 1
# WORLD_SIZE由torch.distributed.launch.py产生，具体数值为 nproc_per_node*node(主机数，这里为1), opt.world_size指进程总数，在这里就是我们使用的卡数
# rank指进程序号，local_rank指本地序号，两者的区别在于前者用于进程间通讯，后者用于本地设备分配,
opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1   # -1 global_rank:进程编号

set_logging(opt.global_rank)
# if opt.global_rank in [-1, 0]:
#     check_git_status()  # 检查你的代码版本是否为最新的(不适用于windows系统),这两行代码最好注释掉

# Resume
if opt.resume:  # resume an interrupted run
    # 如果resume是str,则表示传入的是模型的路径地址
    # get_latest_run()函数获取runs文件夹中最近的last.pt

# DDP mode　多卡训练
if opt.local_rank != -1:

if opt.global_rank in [-1, 0]:          # 创建tensorboard
    logger.info('Start Tensorboard with "tensorboard --logdir %s", view at http://localhost:6006/' % opt.logdir)　
     # 直接点击该网址即可访问tensorboard收集的各种参数曲线
    tb_writer = SummaryWriter(log_dir=log_dir)  # runs/exp0

# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
# 超参数进化列表,括号里分别为(突变规模, 最小值,最大值)

超参进化，看不懂，跳过

2、def train(hyp, opt, device, tb_writer=None)

logger.info(f'Hyperparameters {hyp}')
# 获取记录训练日志的路径
"""
训练日志包括：权重、tensorboard文件、超参数hyp、设置的训练参数opt(也就是epochs,batch_size等),result.txt
result.txt包括: 占GPU内存、训练集的GIOU loss, objectness loss, classification loss, 总loss, 
targets的数量, 输入图片分辨率, 准确率TP/(TP+FP),召回率TP/P ; 
测试集的mAP50, [email protected]:0.95, GIOU loss, objectness loss, classification loss.
还会保存batch<3的ground truth
"""
# 如果设置进化算法则不会传入tb_writer(则为None)，设置一个evolve文件夹作为日志目录
log_dir = Path(tb_writer.log_dir) if tb_writer else Path(opt.logdir) / 'evolve'  # logging directory

# 获取轮次、批次、总批次(涉及到分布式训练)、权重、进程序号(主要用于分布式训练)
epochs, batch_size, total_batch_size, weights, rank = \
        opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank

init_seeds(2 + rank)        # 设置随机种子  rank = -1
# 在神经网络中，参数默认是进行随机初始化的。如果不设置的话每次训练时的初始化都是随机的，
# 导致结果不确定。如果设置初始化，则每次初始化都是固定的

torch.manual_seed(seed)   # 为CPU设置种子用于生成随机数，以使得结果是确定的
# 如果使用多个GPU，应该使用torch.cuda.manual_seed_all()为所有的GPU设置种子

with torch_distributed_zero_first(rank):  # 分布式训练,同步所有进程
    check_dataset(data_dict)  # check  检查下载数据集

state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
model.load_state_dict(state_dict, strict=False)  # load
# 显示加载预训练权重的的键值对和创建模型的键值对
# 如果设置了resume，则会少加载两个键值对(anchors,anchor_grid)

# Freeze
"""
冻结模型层,设置冻结层名字即可,具体可以查看https://github.com/ultralytics/yolov5/issues/679
但作者不鼓励冻结层,因为他的实验当中显示冻结层不能获得更好的性能,参照:https://github.com/ultralytics/yolov5/pull/707
并且作者为了使得优化参数分组可以正常进行,在下面将所有参数的requires_grad设为了True, 其实这里只是给一个freeze的示例
"""
freeze = ['', ]  # parameter names to freeze (full or partial)

# Optimizer
"""
nbs为模拟的batch_size; 就比如默认的话上面设置的opt.batch_size为16,这个nbs就为64，
也就是模型梯度累积了64/16=4(accumulate)次之后,再更新一次模型，变相的扩大了batch_size
"""
nbs = 64  # nominal batch size
accumulate = max(round(nbs / total_batch_size), 1)  # accumulate loss before optimizing
# 根据accumulate设置权重衰减系数
hyp['weight_decay'] *= total_batch_size * accumulate / nbs  # scale weight_decay

pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
# 将模型分成三组(weight、bn, bias, 其他所有参数)优化

# 设置学习率衰减，这里为余弦退火方式进行衰减, 就是根据以下公式lf,epoch和超参数hyp['lrf']进行衰减
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf']  # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)

# Resume
# 初始化开始训练的epoch和最好的结果
# best_fitness是以[0.0, 0.0, 0.1, 0.9]为系数并乘以[精确度, 召回率, [email protected], [email protected]:0.95]再求和所得
# 根据best_fitness来保存best.pt
start_epoch, best_fitness = 0, 0.0

# DP mode
# 分布式训练,参照:https://github.com/ultralytics/yolov5/issues/475
# DataParallel模式,仅支持单机多卡
# rank为进程编号, 这里应该设置为rank=-1则使用DataParallel模式
# rank=-1且gpu数量=1时,不会进行分布式
if cuda and rank == -1 and torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)

# SyncBatchNorm  使用跨卡同步BN
if opt.sync_bn and cuda and rank != -1:
    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
    logger.info('Using SyncBatchNorm()')

# Exponential moving average  为模型创建EMA指数滑动平均,如果GPU进程数大于1,则不创建
ema = ModelEMA(model) if rank in [-1, 0] else None

# DDP mode
# 如果rank不等于-1,则使用DistributedDataParallel模式
# local_rank为gpu编号,rank为进程,例如rank=3，local_rank=0 表示第 3 个进程内的第 1 块 GPU。
if cuda and rank != -1:
    model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)

# Process 0
if rank in [-1, 0]:
    # 更新ema模型的updates参数,保持ema的平滑性
    ema.updates = start_epoch * nb // accumulate  # set EMA updates

# Anchors
"""
计算默认锚点anchor与数据集标签框的长宽比值
标签的长h宽w与anchor的长h_a宽w_a的比值, 即h/h_a, w/w_a都要在(1/hyp['anchor_t'], hyp['anchor_t'])是可以接受的
如果标签框满足上面条件的数量小于总数的99%，则根据k-mean算法聚类新的锚点anchor
"""
if not opt.noautoanchor:
    check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)

"""
设置学习率衰减所进行到的轮次，
目的是打断训练后，--resume接着训练也能正常的衔接之前的训练进行学习率衰减
"""
scheduler.last_epoch = start_epoch - 1  # do not move
# 通过torch1.6自带的api设置混合精度训练
scaler = amp.GradScaler(enabled=cuda)

# Update image weights (optional)
if opt.image_weights:
    # Generate indices
    """
    如果设置进行图片采样策略，
    则根据前面初始化的图片采样权重model.class_weights以及maps配合每张图片包含的类别数
    通过random.choices生成图片索引indices从而进行采样
    """

# Warmup
"""
热身训练(前nw次迭代)
在前nw次迭代中，根据以下方式选取accumulate和学习率
"""
if ni <= nw:
    xi = [0, nw]  # x interp
    # model.gr = np.interp(ni, xi, [0.0, 1.0])  # giou loss ratio (obj_loss = 1.0 or giou)
    accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
    for j, x in enumerate(optimizer.param_groups):
        # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
        """
        bias的学习率从0.1下降到基准学习率lr*lf(epoch)，
        其他的参数学习率从0增加到lr*lf(epoch).
        lf为上面设置的余弦退火的衰减函数
        """
        x['lr'] = np.interp(ni, xi, [0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
        # 动量momentum也从0.9慢慢变到hyp['momentum'](default=0.937)
        if 'momentum' in x:
            x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']])

# Multi-scale  设置多尺度训练，从imgsz * 0.5, imgsz * 1.5 + gs随机选取尺寸
if opt.multi_scale:
    sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size
    sf = sz / max(imgs.shape[2:])  # scale factor
    if sf != 1:
        ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
        imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)

# Forward
with amp.autocast(enabled=cuda):   # 混合精度
    pred = model(imgs)  # forward
    # 计算损失，包括分类损失，objectness损失，框的回归损失
    # loss为总损失值，loss_items为一个元组，包含分类损失，objectness损失，框的回归损失和总损失
    loss, loss_items = compute_loss(pred, targets.to(device), model)  # loss scaled by batch_size
    if rank != -1:
        # 平均不同gpu之间的梯度
        loss *= opt.world_size  # gradient averaged between devices in DDP mode

 # Optimize
 # 模型反向传播accumulate次之后再根据累积的梯度更新一次参数
 if ni % accumulate == 0:
     scaler.step(optimizer)  # optimizer.step
     scaler.update()
     optimizer.zero_grad()
     if ema:
         ema.update(model)

# Scheduler   进行学习率衰减
lr = [x['lr'] for x in optimizer.param_groups]  # for tensorboard
scheduler.step()

# Strip optimizers
"""
模型训练完后，strip_optimizer函数将optimizer从ckpt中去除；
并且对模型进行model.half(), 将Float32的模型->Float16，
可以减少模型大小，提高inference速度
"""

后面再看模型和训练时读取数据的代码