您的位置:

RetinaNet网络结构详解

RetinaNet是Focal Loss for Dense Object Detection这篇论文中提出的一种目标检测网络结构,该网络结构在相同的精度情况下提高了训练速度。RetinaNet基于Focal Loss分类器来加强正样本和负样本之间的区分度,同时引入了Focal Loss检测器来提高检测器的灵敏度。以下是RetinaNet网络结构的详细解析。

一、Anchor-based检测器

RetinaNet的目标检测器是一种Anchor-based检测器,其中Anchor是指在输入图像中的一组预先定义的框(或称锚定框),每个框都是有关尺度和长宽比的离散集合。框与像素之间的映射是通过网络中最后一个卷积层完成的,它将卷积层的特征图与原始输入图像之间生成了一个映射。检测器针对每个Anchor框执行了两个任务:首先,预测所属类别的概率值;其次,预测框向真实边界框的偏移量。在训练期间,对于每个Anchor框,如果预测结果与真实框匹配,则该Anchor框被视为正样本,否则该Anchor框被视为负样本。这种Anchor-based方法使模型可以对不同数量和尺度的物体进行识别和分割。 下面是RetinaNet的Anchor-based检测器的代码实现:
class RetinaNet(nn.Module):
    def __init__(self):
        super(RetinaNet, self).__init__()

        self.fpn = FPN()
        self.cls_head = ClsHead()
        self.reg_head = RegHead()

    def forward(self, x):
        out = self.fpn(x)
        cls_out = []
        reg_out = []

        for feature in out:
            cls_out.append(self.cls_head(feature))
            reg_out.append(self.reg_head(feature))

        return tuple(cls_out), tuple(reg_out)

二、Focal Loss分类器

Focal Loss是针对目标检测任务的一种修改后的二分类损失函数,它通过加权函数来缓解分类器在面对大量简单负样本(例如背景)时的鲁棒性问题。具体来说,该权重函数主要是在标准交叉熵损失中引入一个可调参数,该参数控制与正确分类相关的样本的权重值。当$\alpha$=0.5时,该权重函数将标准交叉熵损失还原为通用的交叉熵损失。Focal Loss通过对易分类的样本进行降权来使分类器更加关注难分类的样本。 下面是RetinaNet考虑Focal Loss的分类器的代码实现:
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()

        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, cls_pred, cls_targets):
        pos_inds = cls_targets > 0
        neg_inds = cls_targets == 0

        pos_pred = cls_pred[pos_inds]
        neg_pred = cls_pred[neg_inds]

        pos_loss = -pos_pred.log() * (1 - pos_pred) ** self.gamma * self.alpha
        neg_loss = -neg_pred.log() * (neg_pred) ** self.gamma * (1 - self.alpha)

        if self.reduction == 'mean':
            num_pos = pos_inds.float().sum()
            pos_loss = pos_loss.sum()
            neg_loss = neg_loss.sum()
            loss = (pos_loss + neg_loss) / num_pos.clamp(min=1)
        else:
            loss = pos_loss.sum() + neg_loss.sum()

        return loss

三、Focal Loss检测器

RetinaNet引入了一个新的检测器,称为Focal Loss检测器,该检测器与Focal Loss分类器共同作用。具体来说,RetinaNet的Focal Loss检测器在分类时考虑了Focal Loss,这意味着该检测器在面对难分类样本时会更加关注,而忽略容易分类的样本。 下面是RetinaNet的Focal Loss检测器的代码实现:
class FocalLossDetection(nn.Module):
    def __init__(self, alpha=0.25, gamma=2, reduction='mean'):
        super(FocalLossDetection, self).__init__()

        self.cls_loss = FocalLoss(alpha, gamma, reduction=reduction)
        self.reg_loss = nn.SmoothL1Loss(reduction=reduction)

    def forward(self, cls_out, reg_out, cls_targets, reg_targets):
        cls_losses = []
        reg_losses = []

        for cls_pred, reg_pred, cls_target, reg_target, in zip(cls_out, reg_out, cls_targets, reg_targets):
            pos_inds = cls_target > 0
            num_pos = pos_inds.float().sum()

            cls_loss = self.cls_loss(cls_pred, cls_target)
            reg_loss = self.reg_loss(pos_pred, pos_target, )

            cls_losses.append(cls_loss)
            reg_losses.append(reg_loss)

        cls_loss = sum(cls_losses) / len(cls_losses)
        reg_loss = sum(reg_losses) / len(reg_losses)

        loss = cls_loss + reg_loss

        return loss

四、RetinaNet网络结构整合

最后,我们将RetinaNet网络结构从头到尾地整理一遍。整个网络结构包括了FPN、ClsHead、RegHead、Focal Loss和Smooth L1损失。其中,FPN生成了多个特征层,而ClsHead和RegHead分别预测类别概率和边框偏移。Focal Loss和Smooth L1损失作为网络的训练损失函数。 下面是整合后的RetinaNet网络结构代码实现:
class FPN(nn.Module):
    def __init__(self):
        super(FPN, self).__init__()

        self.conv6 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)
        self.conv7 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)
        self.conv8 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)
        self.conv9 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)

        self.latent3 = nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)
        self.latent4 = nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)
        self.latent5 = nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)

        self.pred3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.pred4 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.pred5 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        conv3, conv4, conv5, conv6, conv7, conv8, conv9 = x

        lat3 = self.latent3(conv3)
        lat4 = self.latent4(conv4)
        lat5 = self.latent5(conv5)

        p5 = self.pred5(lat5)
        p4 = self.pred4(lat4 + F.interpolate(p5, size=lat4.size()[-2:], mode='nearest'))
        p3 = self.pred3(lat3 + F.interpolate(p4, size=lat3.size()[-2:], mode='nearest'))

        p6 = self.conv6(conv6)
        p7 = self.conv7(F.relu(p6))
        p8 = self.conv8(F.relu(p7))
        p9 = self.conv9(F.relu(p8))

        return p3, p4, p5, p6, p7, p8, p9


class ClsHead(nn.Module):
    def __init__(self):
        super(ClsHead, self).__init__()

        self.output = nn.Conv2d(256, 9, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        x = self.output(x)
        x = x.permute(0, 2, 3, 1)
        x = x.reshape(x.shape[0], -1, 1)

        return x


class RegHead(nn.Module):
    def __init__(self):
        super(RegHead, self).__init__()

        self.output = nn.Conv2d(256, 36, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        x = self.output(x)
        x = x.permute(0, 2, 3, 1)
        x = x.reshape(x.shape[0], -1, 4)

        return x


class RetinaNet(nn.Module):
    def __init__(self):
        super(RetinaNet, self).__init__()

        self.fpn = FPN()
        self.cls_head = ClsHead()
        self.reg_head = RegHead()

        self.focal_loss_detection = FocalLossDetection()

    def forward(self, x, cls_targets, reg_targets):
        out = self.fpn(x)
        cls_out = []
        reg_out = []

        for feature in out:
            cls_out.append(self.cls_head(feature))
            reg_out.append(self.reg_head(feature))

        loss = self.focal_loss_detection(cls_out, reg_out, cls_targets, reg_targets)

        return loss, cls_out, reg_out