Yolo v1网络结构

yolo v1的网络结构如下图所示

其输入为$448\times 448\times 3$的RGB图像，输出为$7\times 7\times 30$的预测张量，其中$7\times 7$是空间重复，$30$则包含两个部分，分别是预测框的位置和类别数，$0,1,2,3,4$分别是$x,y,w,h,c$，即坐标宽高和置信度，$5,6,7,8,9$是第二组的$x,y,w,h,c$，$\mathrm{10..}$则是20种类别，如果是当前类别，则置为$1$否则置$0$。

总计24层卷积层，如下表所示

其Pytorch实现为

import torch
import torch.nn as nn
GRID = 7
VOC_CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat',
               'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
NUM_CLASSES = len(VOC_CLASSES)
class YOLOv1(nn.Module):
    def __init__(self, num_classes=NUM_CLASSES):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3), nn.LeakyReLU(0.1), nn.MaxPool2d(2, stride=2),
            nn.Conv2d(64, 192, 3, padding=1), nn.LeakyReLU(0.1), nn.MaxPool2d(2, stride=2),
            nn.Conv2d(192, 128, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(128, 256, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(256, 256, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(256, 512, 3, padding=1), nn.LeakyReLU(0.1),nn.MaxPool2d(2, stride=2),
            nn.Conv2d(512, 256, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(256, 512, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(512, 256, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(256, 512, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(512, 256, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(256, 512, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(512, 256, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(256, 512, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(512, 512, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(512, 1024, 3, padding=1), nn.LeakyReLU(0.1), nn.MaxPool2d(2, stride=2),
            nn.Conv2d(1024, 512, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(512, 1024, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(1024, 512, 1), nn.LeakyReLU(0.1), 
            nn.Conv2d(512, 1024, 3, padding=1), nn.LeakyReLU(0.1),
            nn.Conv2d(1024, 1024, 3, padding=1), nn.LeakyReLU(0.1), 
            nn.Conv2d(1024, 1024, 3, padding=1), nn.LeakyReLU(0.1)
        )
        self.downsample = nn.AvgPool2d(2, stride=2)
        self.fc = nn.Sequential(
            nn.Linear(GRID**2 * 1024, 4096), nn.LeakyReLU(0.1), nn.Dropout(0.5),
            nn.Linear(4096, GRID**2 * (num_classes + 10))
        )

    def forward(self, x):
        x = self.conv(x)
        x = self.downsample(x)
        x = x.view(x.size(0), -1)
        out = self.fc(x).view(-1, GRID, GRID, NUM_CLASSES + 10)
        return out

损失函数

Yolo v1的损失函数为

$$\begin{array}{rl}L=& {\lambda }_c\sum _{i=0}^{S^2}\sum _{j=0}^B{\mathrm{𝟙}}_{ij}\left[(x_i-{\hat{x}}_i{)}^2+(y_i-{\hat{y}}_i{)}^2+{\left(\sqrt{{\omega }_i}-\sqrt{{\hat{\omega }}_i}\right)}^2+{\left(\sqrt{h_i}-\sqrt{{\hat{h}}_i}\right)}^2\right]\\ & +\sum _{i=0}^{S^2}\sum _{j=0}^B{\mathrm{𝟙}}_{ij}(C_i-{\hat{C}}_i{)}^2+{\lambda }_n\sum _{i=0}^{S^2}\sum _{j=0}^B{\mathrm{𝟘}}_{ij}(C_i-{\hat{C}}_j{)}^2\\ & +\sum _{i=0}^{S^2}{\mathrm{𝟙}}_{ij}\sum _{c\in classes}{\left[p_i(c)-{\hat{p}}_i(c)\right]}^2\end{array}$$

式中，

• $S=7,B=2$分别是网格划分数和每格预测框数
• ${\mathrm{𝟙}}_{ij}$(原文中是${\mathrm{𝟙}}_{ij}^{obj}$)表示第$i$格第$j$框负责预测物体时为1，否则为0
• ${\mathrm{𝟘}}_{ij}$(原文中是${\mathrm{𝟙}}_{ij}^{noobj}$)表示第$i$格第$j$框不负责预测物体时为1，否则为0
• ${\lambda }_c=0.5,{\lambda }_n=0.5$（原文${\lambda }_{coord},{\lambda }_{noobj}$）
• $\hat{x},\hat{y},\hat{w},\hat{h}$ 为预测值
• $C_i$为预测置信度；${\hat{C}}_i$为真实置信度
• $p_i(c),{\hat{p}}_i(c)$为预测类别概率和真实类别

这三个组成部分分别表示

• 坐标定位损失
• 物体置信度损失
• 类别分类损失

在具体实现时，要注意$w,h$需要开根号，而$x,y$不需要

sqrt = lambda x: torch.sqrt(torch.clamp(x, min=1e-6))
class YOLOLoss(nn.Module):
    def __init__(self, c=5.0, n=0.5):
        super().__init__()
        self.c, self.n = c, n

    def forward(self, pred, target, mask):
        xyp1, whp1, cp1, xyp2, whp2, cp2, clsp = torch.split(
            pred, [2, 2, 1, 2, 2, 1, NUM_CLASSES], dim=-1)
        xyt1, wht1, ct1, xyt2, wht2, ct2, clst = torch.split(
            target, [2, 2, 1, 2, 2, 1, NUM_CLASSES], dim=-1)
        m1, m2 = mask[..., 0:1], mask[..., 1:2]
        cell_mask = torch.clamp(m1 + m2, max=1.0)
        
        dxy1 = (xyp1 - xyt1) ** 2
        dwh1 = (sqrt(whp1) - sqrt(wht1)) ** 2
        dxy2 = (xyp2 - xyt2) ** 2
        dwh2 = (sqrt(whp2) - sqrt(wht2)) ** 2
        coordL = self.c * torch.sum(m1*(dxy1 + dwh1) + m2*(dxy2 + dwh2))
        confL = torch.sum(m1*(cp1-ct1)**2 + m2*(cp2-ct2)**2)
        confL += self.n * torch.sum((1-m1)*cp1**2 + (1-m2)*cp2**2)
        clsL = torch.sum(cell_mask * (clsp - clst)**2)
        return (coordL + confL + clsL)/pred.size(0)