YOLOv3-pytorch代码详解

YOLOV3的YOLOLayer代码详解

YOLOV3的网络结构网上已经有很多人写了很多高质量的文章，我不打算再重复这一项工作，这里仅转载一下这张结构图：

转载于：江大白

这张图很清晰的展示了YOLOV3的网络结构，不过需要注意的是，YOLOV3允许输入不同大小的图片，可以是416也可以是608，不同大小的图片在网络输出时的三个尺度的feature尺度也不同。

backbone为Darknet53，Neck部分为标准FPN，输出为三个尺度的预测。简单提一下与YOLOV4在网络结构上的区别：backbone为CSPDarkNet53，Neck部分为标准FPN+PAN结构，即上采样融合之后再下采样融合，输出与YOLOV3相同。当然，YOLOV4包含了巨量的trick，区别完全不止这一点，不过这个文章主要讲YOLOV3，就不多做展开。

下面开始手撕YOLOLayer。

因为比较懒，就不再造轮子，引用一下别人的图：

转载于：https://www.cnblogs.com/ywheunji/p/10809695.html

这张图与YOLOV3结构图输入大小不同，不过这个都是允许的。这张图很清晰的描述了YOLOV3输出部分的内容。以第一个尺寸（13*13）为例，这条长柱横截面为（13*13），即为feature的像素，然后长柱被分为三截，这三截即这个size的feature分配到的三个尺寸的anchor（YOLOV3通过对coco数据集的Bbox的size聚类，得到九类size的anchor，这九类被分配到三个尺寸的输出feature上）。这三截的长度为（4+1+类别数量），这里的4为Bbox的中心坐标（x,y）以及长宽（w,h），1为这个anchor包含目标的置信度，类别数量为这个anchor的label，比如Coco数据集的类别数为80。

清楚了网络的输出结构，就开始解释先验框与目标框的匹配策略，生成网络输出目标，以及计算实际网络输出与这个目标之间的loss了。

上代码：

class YOLOLayer(nn.Module):
    """Detection layer"""

    def __init__(self, anchors, num_classes, img_dim=416):
        super(YOLOLayer, self).__init__()
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.ignore_thres = 0.5
        self.mse_loss = nn.MSELoss()
        self.bce_loss = nn.BCELoss()
        self.obj_scale = 1
        self.noobj_scale = 100
        self.metrics = {}
        self.img_dim = img_dim
        self.grid_size = 0  # grid size

    def compute_grid_offsets(self, grid_size, cuda=True):
        self.grid_size = grid_size
        g = self.grid_size
        FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
        self.stride = self.img_dim / self.grid_size
        # Calculate offsets for each grid
        self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor)
        self.grid_y = torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor)
        self.scaled_anchors = FloatTensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors])
        self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))
        self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))

    def forward(self, x, targets=None, img_dim=None):

        # Tensors for cuda support
        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
        ByteTensor = torch.cuda.ByteTensor if x.is_cuda else torch.ByteTensor

        self.img_dim = img_dim
        num_samples = x.size(0)
        grid_size = x.size(2)

        prediction = (
            x.view(num_samples, self.num_anchors, self.num_classes + 5, grid_size, grid_size)
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )

        # Get outputs
        x = torch.sigmoid(prediction[..., 0])  # Center x
        y = torch.sigmoid(prediction[..., 1])  # Center y
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height
        # 每个预测框包含目标的概率
        pred_conf = torch.sigmoid(prediction[..., 4])  # Conf
        # 每个预测框属于每个类别的概率，注意，这里使用的sigmod而非softmax，所以超过阈值的可能有多个，即anchor可以属于多个类
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

        # If grid size does not match current we compute new offsets
        # 因为允许不同size的图片输入，有可能上一次的图片和本次的图片size不同，那么就需要更新grid_x、grid_y、anchor_w、anchor_h
        if grid_size != self.grid_size:
            self.compute_grid_offsets(grid_size, cuda=x.is_cuda)

        # Add offset and scale with anchors
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + self.grid_x
        pred_boxes[..., 1] = y.data + self.grid_y
        pred_boxes[..., 2] = torch.exp(w.data) * self.anchor_w
        pred_boxes[..., 3] = torch.exp(h.data) * self.anchor_h

        # 这里的output输出维度为 bach_size 3*gred_size*grid_size 5+num_cls
        # 中间的维度即为每张图片的所有anchor数，最后一个维度即为每个anchor的输出（x, y, w, h, c, p）
        # 其中 c 为anchor包含目标的置信度， p 为每个anchor属于每个类别的概率
        output = torch.cat(
            (
                pred_boxes.view(num_samples, -1, 4) * self.stride,
                pred_conf.view(num_samples, -1, 1),
                pred_cls.view(num_samples, -1, self.num_classes),
            ),
            -1,
        )
        # 如果没有targets，则表示本次为推理，不需要进行训练，直接返回结果
        if targets is None:
            return output, 0
        else:
            iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf = build_targets(
                pred_boxes=pred_boxes,
                pred_cls=pred_cls,
                target=targets,
                anchors=self.scaled_anchors,
                ignore_thres=self.ignore_thres,
            )

            obj_mask = obj_mask.bool()  # convert int8 to bool
            noobj_mask = noobj_mask.bool()  # convert int8 to bool

            # Loss : Mask outputs to ignore non-existing objects (except with conf. loss)
            # 注意，这里计算x,y,w,h的loss时都只计算有目标的anchor（即与gt框最匹配的anchor，有多少个gt框则有多少个anchor）
            loss_x = self.mse_loss(x[obj_mask], tx[obj_mask])
            loss_y = self.mse_loss(y[obj_mask], ty[obj_mask])
            loss_w = self.mse_loss(w[obj_mask], tw[obj_mask])
            loss_h = self.mse_loss(h[obj_mask], th[obj_mask])
            # 计算包含目标的置信度loss需要计算包含目标和没有包含目标的anchor的loss，但是需要注意的是
            # 这里的包含目标为 与gt框最匹配的anchor，不论其与gt框的IOU是否低于阈值
            # 而不含目标为 除去与gt框最匹配的anchor和与gt框IOU超过阈值的anchor后的所有anchor，
            # 所以IOU超过阈值但并非与gt最匹配的anchor，在计算置信度损失时将会被忽略
            loss_conf_obj = self.bce_loss(pred_conf[obj_mask], tconf[obj_mask])
            loss_conf_noobj = self.bce_loss(pred_conf[noobj_mask], tconf[noobj_mask])
            loss_conf = self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj
            # pred_cls 为激活函数sigmod的输出，tcls为gt的真实标签的one-hot编码，分类损失为两者的交叉熵损失
            loss_cls = self.bce_loss(pred_cls[obj_mask], tcls[obj_mask])
            total_loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls

            # Metrics
            cls_acc = 100 * class_mask[obj_mask].mean()
            conf_obj = pred_conf[obj_mask].mean()
            conf_noobj = pred_conf[noobj_mask].mean()
            conf50 = (pred_conf > 0.5).float()
            iou50 = (iou_scores > 0.5).float()
            iou75 = (iou_scores > 0.75).float()
            detected_mask = conf50 * class_mask * tconf
            precision = torch.sum(iou50 * detected_mask) / (conf50.sum() + 1e-16)
            recall50 = torch.sum(iou50 * detected_mask) / (obj_mask.sum() + 1e-16)
            recall75 = torch.sum(iou75 * detected_mask) / (obj_mask.sum() + 1e-16)

            self.metrics = {
                "loss": to_cpu(total_loss).item(),
                "x": to_cpu(loss_x).item(),
                "y": to_cpu(loss_y).item(),
                "w": to_cpu(loss_w).item(),
                "h": to_cpu(loss_h).item(),
                "conf": to_cpu(loss_conf).item(),
                "cls": to_cpu(loss_cls).item(),
                "cls_acc": to_cpu(cls_acc).item(),
                "recall50": to_cpu(recall50).item(),
                "recall75": to_cpu(recall75).item(),
                "precision": to_cpu(precision).item(),
                "conf_obj": to_cpu(conf_obj).item(),
                "conf_noobj": to_cpu(conf_noobj).item(),
                "grid_size": grid_size,
            }

            return output, total_loss

这里主要看forward函数。进来首先获取到bach_size以及grid_size，如果输入是416，那么这个grid_size即为13，然后将输入变形为（bach_size, 3, grid_size, grid_size, num_classes+5）。接下来遇到第一个需要注意的点，计算网络输出，x、y、w、h以及pred_conf都是常规输出，但是pred_cls也是用sigmoid作为激活函数，这样就与SSD有了对比。SSD是以softmax作为最后的激活函数，因此一个anchor只能对应单一的一个标签，而以sigmoid作为激活函数，则可能有多个预测值超过阈值，即一个anchor可能对应于多个标签，这是很有益的，比如当标签中含有person和man时，一个男人就属于这两个标签，SSD只能在这两个标签中选一个，即预测这个男人是一个人或者是个男人，而YOLOV3则可以输出两个标签，即预测这个男人是一个人而且是一个男人。

然后更新一些后面需要的参数，比如grid_x、grid_y、anchor_w、anchor_h。注意下面两行代码：

pred_boxes[..., 0] = x.data + self.grid_x
pred_boxes[..., 1] = y.data + self.grid_y

grid_x、grid_y为整数，且两者组合代表了feature的像素坐标，然后以x、y分别对应相加，即是预测的Bbox中心位置的准确坐标。由此可以看出，输出的其实不是绝对的坐标，而是对应于自自身坐标的偏移。这样说起来可能比较抽象，那么就简单的模拟一下x、y的值，然后看其对应的信息：

这里为了好看，将x、y的size都设为5*5，实际上如果输入为416，其size应该是13*13，不过意思都一样，不要纠结这种小细节。以x[2,3]与y[2,3]为例，其负责预测的feature的（2,3）坐标的目标，因为其值为（0.5755,0.5782），这就意味着其预测的实际的目标中心应该在（2.5755,3.5782）.

下面的代码我都已经注释得很清楚了，不过这里还需要跳入build_targets函数，这个函数即为根据实际的目标构建本次训练的目标，也就是用ground truth构建出希望网络输出的值。

下面上代码：

def build_targets(pred_boxes, pred_cls, target, anchors, ignore_thres):

    ByteTensor = torch.cuda.ByteTensor if pred_boxes.is_cuda else torch.ByteTensor
    FloatTensor = torch.cuda.FloatTensor if pred_boxes.is_cuda else torch.FloatTensor

    nB = pred_boxes.size(0)  # bach_size
    nA = pred_boxes.size(1)  # anchor数目
    nC = pred_cls.size(-1)   # 类别数目
    nG = pred_boxes.size(2)  # grid_size

    # Output tensors
    obj_mask = ByteTensor(nB, nA, nG, nG).fill_(0)
    noobj_mask = ByteTensor(nB, nA, nG, nG).fill_(1)
    class_mask = FloatTensor(nB, nA, nG, nG).fill_(0)
    iou_scores = FloatTensor(nB, nA, nG, nG).fill_(0)
    tx = FloatTensor(nB, nA, nG, nG).fill_(0)
    ty = FloatTensor(nB, nA, nG, nG).fill_(0)
    tw = FloatTensor(nB, nA, nG, nG).fill_(0)
    th = FloatTensor(nB, nA, nG, nG).fill_(0)
    tcls = FloatTensor(nB, nA, nG, nG, nC).fill_(0)

    # Convert to position relative to box
    target_boxes = target[:, 2:6] * nG
    gxy = target_boxes[:, :2]
    gwh = target_boxes[:, 2:]
    # Get anchors with best iou
    # 这里需要注意，和SSD不同，这里是只根据Bbox形状做IOU，没有坐标，可以理解为这一步只给出每种anchor与gt的相似度
    #### 维度为 3 × target_num,因此，iou.max(0)返回的分别为每个target与最匹配的anchor的IOU以及这个最匹配的anchor的index
    ious = torch.stack([bbox_wh_iou(anchor, gwh) for anchor in anchors])
    best_ious, best_n = ious.max(0)
    # Separate target values
    # 注意，这里的target第一个维度为目标所属的bach的index，因此b代表目标所属的bach index
    b, target_labels = target[:, :2].long().t()
    gx, gy = gxy.t()
    gw, gh = gwh.t()
    gi, gj = gxy.long().t()   # gxy.long()将坐标转换成整数，然后转置，将x,y分开，获得gt的中心坐标gi，gj
    # mask的维度分别为 bach_size  anchor_num  grid_size  grid_size，前面已经求到了与gt框最相似的anchor的index
    # 因此在gi，gj坐标第index个anchor必然与gt最匹配，在最匹配的位置就认为是有目标的，不管IOU是否超过阈值
    obj_mask[b, best_n, gj, gi] = 1
    noobj_mask[b, best_n, gj, gi] = 0

    # 如果某个anchor与gt框IOU大于阈值，则不认为这里没有目标（注意，也没有认为这里有目标）
    for i, anchor_ious in enumerate(ious.t()):
        noobj_mask[b[i], anchor_ious > ignore_thres, gj[i], gi[i]] = 0

    # 这里的tx,ty是gt框中心相对于这个负责预测它的anchor的中心点坐标的小数点偏移，并非是直接的x,y坐标
    tx[b, best_n, gj, gi] = gx - gx.floor()
    ty[b, best_n, gj, gi] = gy - gy.floor()
    # 这里的宽高是gt的宽高与负责预测它的anchor的宽高的比例
    tw[b, best_n, gj, gi] = torch.log(gw / anchors[best_n][:, 0] + 1e-16)
    th[b, best_n, gj, gi] = torch.log(gh / anchors[best_n][:, 1] + 1e-16)
    # One-hot encoding 注意target_labels可能是多个，比如是person又是man，即多标签
    tcls[b, best_n, gj, gi, target_labels] = 1
    # Compute label correctness and iou at best anchor
    class_mask[b, best_n, gj, gi] = (pred_cls[b, best_n, gj, gi].argmax(-1) == target_labels).float()
    iou_scores[b, best_n, gj, gi] = bbox_iou(pred_boxes[b, best_n, gj, gi], target_boxes, x1y1x2y2=False)

    tconf = obj_mask.float()
    return iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf

这里需要注意有两行代码，与上面我说的x、y的预测为其预测目标相对于其本身坐标的偏移这一点对应：

# 这里的tx,ty是gt框中心相对于这个负责预测它的anchor的中心点坐标的小数点偏移，并非是直接的x,y坐标
tx[b, best_n, gj, gi] = gx - gx.floor()
ty[b, best_n, gj, gi] = gy - gy.floor()

相信只要理解了这两段代码，对于YOLOV3的核心思想就有了一个清晰的认识。至于其他的网络结构都比较简单，随便打印一下都可以看懂。

这些我都是经过训练Coco数据集以及自己的数据集调试之后理解出来的，应该没什么问题，如果有问题欢迎提出异议，以期“越辩越明”。