前言
看了 Yolov3 的論文之后，發現這論文寫的真的是很簡短，神經網絡的具體結構和損失函數的公式都沒有給出。所以這里參考了許多前人的博客和代碼，下面進入正題。
網絡結構
Yolov3 將主干網絡換成了 darknet53，整體的網絡結構如下圖所示（圖片來自【論文解讀】Yolo三部曲解讀——Yolov3）：
這里的 CONV 具體結構是 1 個 Conv2d + 1 個 BatchNorm2d + 1個 LeakyReLU （除了 Feature Map 1、2、3 前的 1×1 CONV），代碼如下：
class ConvBlock(nn.Module):
""" 卷積塊 """
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size,
stride, padding, bias=False),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.1)
)
def forward(self, x):
return self.features(x)
ResBlock 里面進行了兩次卷積，最后還有一個跳連接，代碼如下：
class ResidualUnit(nn.Module):
""" 殘差單元 """
def __init__(self, in_channels: int):
super().__init__()
self.features = nn.Sequential(
ConvBlock(in_channels, in_channels//2, 1, padding=0),
ConvBlock(in_channels//2, in_channels, 3),
)
def forward(self, x):
y = self.features(x)
return x+y
ResOperator(n) 包括一個卷積核大小為 3×3、步長為 2 的 Convolutional 和 n 個 ResBlock，代碼如下：
class ResidualBlock(nn.Module):
""" 殘差塊 """
def __init__(self, in_channels: int, n_residuals=1):
"""
Parameters
----------
in_channels: int
輸入通道數
n_residuals: int
殘差單元的個數
"""
super().__init__()
self.conv = ConvBlock(in_channels, in_channels*2, 3, stride=2)
self.residual_units = nn.Sequential(*[
ResidualUnit(2*in_channels) for _ in range(n_residuals)
])
def forward(self, x):
return self.residual_units(self.conv(x))
在 Darknet53 中共有 1+2+8+8+4=23 個 ResOperator，假設輸入 Darknet53 的圖像維度為 256×256×3，那么最后會輸出 3 個特征圖，維度分別是 32×32×256、16×16×512 和 8×8×1024，大小分別是原始圖像的 1/8、1/16 和 1/32（說明輸入 Darknet 的圖像的寬高需要是 32 的倍數），這些特征圖會在后面接着進行卷積。
class Darknet(nn.Module):
""" 主干網絡 """
def __init__(self):
super().__init__()
self.conv = ConvBlock(3, 32, 3)
self.residuals = nn.ModuleList([
ResidualBlock(32, 1),
ResidualBlock(64, 2),
ResidualBlock(128, 8),
ResidualBlock(256, 8),
ResidualBlock(512, 4),
])
def forward(self, x):
"""
Parameters
----------
x: Tensor of shape `(N, 3, H, W)`
輸入圖像
Returns
-------
x1: Tensor of shape `(N, 1024, H/32, W/32)`
x2: Tensor of shape `(N, 512, H/16, W/16)`
x3: Tensor of shape `(N, 256, H/8, W/8)`
"""
x3 = self.conv(x)
for layer in self.residuals[:-2]:
x3 = layer(x3)
x2 = self.residuals[-2](x3)
x1 = self.residuals[-1](x2)
return x1, x2, x3
8×8×1024 特征圖會先經過 YoloBlock，得到 8×8×512 的特征圖 x1，YoloBlock 的代碼如下所示：
class YoloBlock(nn.Module):
""" Yolo 塊 """
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.features = nn.Sequential(*[
ConvBlock(in_channels, out_channels, 1, padding=0),
ConvBlock(out_channels, out_channels*2, 3, padding=1),
ConvBlock(out_channels*2, out_channels, 1, padding=0),
ConvBlock(out_channels, out_channels*2, 3, padding=1),
ConvBlock(out_channels*2, out_channels, 1, padding=0),
])
def forward(self, x):
return self.features(x)
x1 在經過后續的兩次卷積后，得到維度為 8×8×255 的第一個特征圖 y1，后面會解釋特征圖的含義。
為了融合 x1 中所包含的高級語義， Yolov3 先 x1 進行了 1×1 大小的卷積，減少通道數，接着進行上采樣，使特征圖的維度變為 16×16×256，再與 darknet53 輸出的 16×16×512 的特征圖沿着通道方向進行連結，得到 16×16×768 的特征圖。該特征圖經過 YoloBlock 和后續兩次卷積后得到維度為 16×16×255 的第二個特征圖 y2。y3 同理。總結下來整個 Yolov3 的代碼為：
class Yolo(nn.Module):
""" Yolo 神經網絡 """
def __init__(self, n_classes: int, image_size=416, anchors: list = None, nms_thresh=0.45):
"""
Parameters
----------
n_classes: int
類別數
image_size: int
圖片尺寸，必須是 32 的倍數
anchors: list
輸入圖像大小為 416 時對應的先驗框
nms_thresh: float
非極大值抑制的交並比閾值，值越大保留的預測框越多
"""
super().__init__()
if image_size <= 0 or image_size % 32 != 0:
raise ValueError("image_size 必須是 32 的倍數")
# 先驗框
anchors = anchors or [
[[116, 90], [156, 198], [373, 326]],
[[30, 61], [62, 45], [59, 119]],
[[10, 13], [16, 30], [33, 23]]
]
anchors = np.array(anchors, dtype=np.float32)
anchors = anchors*image_size/416
self.anchors = anchors.tolist()
self.n_classes = n_classes
self.image_size = image_size
self.darknet = Darknet()
self.yolo1 = YoloBlock(1024, 512)
self.yolo2 = YoloBlock(768, 256)
self.yolo3 = YoloBlock(384, 128)
# YoloBlock 后面的卷積部分
out_channels = (n_classes+5)*3
self.conv1 = nn.Sequential(*[
ConvBlock(512, 1024, 3),
nn.Conv2d(1024, out_channels, 1)
])
self.conv2 = nn.Sequential(*[
ConvBlock(256, 512, 3),
nn.Conv2d(512, out_channels, 1)
])
self.conv3 = nn.Sequential(*[
ConvBlock(128, 256, 3),
nn.Conv2d(256, out_channels, 1)
])
# 上采樣
self.upsample1 = nn.Sequential(*[
nn.Conv2d(512, 256, 1),
nn.Upsample(scale_factor=2)
])
self.upsample2 = nn.Sequential(*[
nn.Conv2d(256, 128, 1),
nn.Upsample(scale_factor=2)
])
# 探測器，用於處理輸出的特征圖，后面會講到
self.detector = Detector(
self.anchors, image_size, n_classes, conf_thresh=0.1, nms_thresh=nms_thresh)
def forward(self, x):
"""
Parameters
----------
x: Tensor of shape `(N, 3, H, W)`
輸入圖像
Returns
-------
y1: Tensor of shape `(N, 255, H/32, W/32)`
最小的特征圖
y2: Tensor of shape `(N, 255, H/16, W/16)`
中等特征圖
y3: Tensor of shape `(N, 255, H/8, W/8)`
最大的特征圖
"""
x1, x2, x3 = self.darknet(x)
x1 = self.yolo1(x1)
y1 = self.conv1(x1)
x2 = self.yolo2(torch.cat([self.upsample1(x1), x2], 1))
y2 = self.conv2(x2)
x3 = self.yolo3(torch.cat([self.upsample2(x2), x3], 1))
y3 = self.conv3(x3)
return y1, y2, y3
先驗框
k-means 算法
在介紹先驗框之前，有必要說下 k-means 聚類算法。假設我們有一個樣本集 \(D=\left\{ \boldsymbol{x_1}, \boldsymbol{x_2}, ..., \boldsymbol{x_m} \right\}\) 包含 \(m\) 個樣本，其中每一個樣本 \(\boldsymbol{x_i}=[x_{i1}, x_{i2}, ..., x_{in}]^T\) 是 \(n\) 維空間中的一個點。k-means 算法會將樣本集划分為 \(k\) 個不相交的簇 \(\left\{ C_l | l=1,2,...,k \right\}\) 來最小化平方誤差：
其中 $\boldsymbol{\mu_j}=\frac{1}{|C_j|}\sum_{\boldsymbol{x}\in C_j} \boldsymbol{x} $ 是簇 \(C_j\) 的均值向量（聚類中心）。簡單來說，上面這個公式刻畫了簇內樣本圍繞簇均值向量的緊密程度，\(E\) 值越小則簇內樣本相似度越高。k-means 算法法采用了貪心策略，通過迭代優化來近似求解 \(\boldsymbol{\mu_j}\)。流程如下：
- 從樣本集 \(D\) 中隨機選擇 \(k\) 個樣本作為初始均值向量 \(\left\{ \boldsymbol{\mu_1},\boldsymbol{\mu_2},...,\boldsymbol{\mu_k} \right\}\)
- 對於每一個樣本 \(\boldsymbol{x_i} (1\le i \le m)\) 計算它與各聚類中心 $\boldsymbol{\mu_j}\left(1\le j \le k \right) $的距離 \(d_{ij}=|| \boldsymbol{x_i} - \boldsymbol{\mu_j} ||_2\)
- 根據距離最近的均值向量確定的簇標記 \(\lambda_{i}={\rm argmin} _{j \in \left\{ 1,2,...,k \right\} d_{ij} }\)，將 \(\boldsymbol{x_i}\) 划入響應的簇 \(C_{\lambda_i}=C_{\lambda_{i} } \cup \{ \boldsymbol{x_i} \}\)
- 重新計算聚類中心 $\boldsymbol{\mu_j'}=\frac{1}{|C_j|}\sum_{\boldsymbol{x}\in C_j} \boldsymbol{x} $
- 如果所有的聚類中心都滿足 \(\boldsymbol{\mu_j'}=\boldsymbol{\mu_j}(1\le j \le k)\)，聚類結束，否則回到步驟 2 接着迭代
下圖展示了在西瓜的含糖量和密度張成的二維空間中， \(k=3\) 時 k-means 算法迭代過程（圖片來自周志強老師的《機器學習》）
Yolov3 中的 k-means
上述的 k-means 算法使用歐式距離作為樣本和聚類中心的距離度量，而 Yolov3 中則換了一種距離度量方式。假設邊界框樣本集為 \(D=\left\{ \boldsymbol{x_1}, \boldsymbol{x_2}, ..., \boldsymbol{x_m} \right\}\) 包含 \(m\) 個邊界框，其中每一個樣本 \(\boldsymbol{x_i}=[x_{i1}, x_{i2}, x_{i3}, x_{i4}]^T=[0, 0, w, h]^T\) 是 \(4\) 維空間中的一個點（已歸一化），這里把所有的邊界框的左上角都平移到了原點處。在 Yolov3 中，樣本 \(\boldsymbol{x_i}\) 和聚類中心 \(\boldsymbol{\mu_j}\) 的距離寫作
可以看到樣本和聚類中心的交並比越大，距離就越小，將這個公式作為距離度量可以說是十分的巧妙。
Yolov3 的作者對 COCO 數據集中的邊界框進行了 k-means 聚類，得到了 9 個先驗框，這些先驗框是歸一化后的先驗框再乘以圖像尺寸 416 得到的，由於神經網絡輸出了 3 個不同大小的特征圖，所以將先驗框分為 3 類：
- 大尺度先驗框：[116, 90], [156, 198], [373, 326]，對應最小的那個特征圖 13×13×255
- 中尺度先驗框：[30, 61], [62, 45], [59, 119]，對應中等大小的特征圖 26×26×255
- 小尺度先驗框：[10, 13], [16, 30], [33, 23]，對應最大的特征圖 52×52×255
在訓練自己的數據集時，COCO數據集的先驗框可能不適用，所以應該對邊界框重新聚類。下面給出聚類的代碼：
# coding:utf-8
import glob
from xml.etree import ElementTree as ET
import numpy as np
def iou(box: np.ndarray, boxes: np.ndarray):
""" 計算一個邊界框和多個邊界框的交並比
Parameters
----------
box: `~np.ndarray` of shape `(4, )`
邊界框
boxes: `~np.ndarray` of shape `(n, 4)`
其他邊界框
Returns
-------
iou: `~np.ndarray` of shape `(n, )`
交並比
"""
# 計算交集
xy_max = np.minimum(boxes[:, 2:], box[2:])
xy_min = np.maximum(boxes[:, :2], box[:2])
inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
inter = inter[:, 0]*inter[:, 1]
# 計算並集
area_boxes = (boxes[:, 2]-boxes[:, 0])*(boxes[:, 3]-boxes[:, 1])
area_box = (box[2]-box[0])*(box[3]-box[1])
retun inter/(area_box+area_boxes-inter)
class AnchorKmeans:
""" 先驗框聚類 """
def __init__(self, annotation_dir: str):
self.annotation_dir = annotation_dir
self.bbox = self.get_bbox()
def get_bbox(self) -> np.ndarray:
""" 獲取所有的邊界框 """
bbox = []
for path in glob.glob(f'{self.annotation_dir}/*xml'):
root = ET.parse(path).getroot()
# 圖像的寬度和高度
w = int(root.find('size/width').text)
h = int(root.find('size/height').text)
# 獲取所有邊界框
for obj in root.iter('object'):
box = obj.find('bndbox')
# 歸一化坐標
xmin = int(box.find('xmin').text)/w
ymin = int(box.find('ymin').text)/h
xmax = int(box.find('xmax').text)/w
ymax = int(box.find('ymax').text)/h
bbox.append([0, 0, xmax-xmin, ymax-ymin])
return np.array(bbox)
def get_cluster(self, n_clusters=9, metric=np.median):
""" 獲取聚類結果
Parameters
----------
n_clusters: int
聚類數
metric: callable
選取聚類中心點的方式
"""
rows = self.bbox.shape[0]
if rows < n_clusters:
raise ValueError("n_clusters 不能大於邊界框樣本數")
last_clusters = np.zeros(rows)
clusters = np.ones((n_clusters, 2))
distances = np.zeros((rows, n_clusters)) # type:np.ndarray
# 隨機選取出幾個點作為聚類中心
# np.random.seed(1)
clusters = self.bbox[np.random.choice(rows, n_clusters, replace=False)]
# 開始聚類
while True:
# 計算距離
distances = 1-self.iou(clusters)
# 將每一個邊界框划到一個聚類中
nearest_clusters = distances.argmin(axis=1)
# 如果聚類中心不再變化就退出
if np.array_equal(nearest_clusters, last_clusters):
break
# 重新選取聚類中心
for i in range(n_clusters):
clusters[i] = metric(self.bbox[nearest_clusters == i], axis=0)
last_clusters = nearest_clusters
return clusters[:, 2:]
def average_iou(self, clusters: np.ndarray):
""" 計算 IOU 均值
Parameters
----------
clusters: `~np.ndarray` of shape `(n_clusters, 2)`
聚類中心
"""
clusters = np.hstack((np.zeros((clusters.shape[0], 2)), clusters))
return np.mean([np.max(iou(bbox, clusters)) for bbox in self.bbox])
def iou(self, clusters: np.ndarray):
""" 計算所有邊界框和所有聚類中心的交並比
Parameters
----------
clusters: `~np.ndarray` of shape `(n_clusters, 4)`
聚類中心
Returns
-------
iou: `~np.ndarray` of shape `(n_bbox, n_clusters)`
交並比
"""
bbox = self.bbox
A = self.bbox.shape[0]
B = clusters.shape[0]
xy_max = np.minimum(bbox[:, np.newaxis, 2:].repeat(B, axis=1),
np.broadcast_to(clusters[:, 2:], (A, B, 2)))
xy_min = np.maximum(bbox[:, np.newaxis, :2].repeat(B, axis=1),
np.broadcast_to(clusters[:, :2], (A, B, 2)))
# 計算交集面積
inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
inter = inter[:, :, 0]*inter[:, :, 1]
# 計算每個矩陣的面積
area_bbox = ((bbox[:, 2]-bbox[:, 0])*(bbox[:, 3] -
bbox[:, 1]))[:, np.newaxis].repeat(B, axis=1)
area_clusters = ((clusters[:, 2] - clusters[:, 0])*(
clusters[:, 3] - clusters[:, 1]))[np.newaxis, :].repeat(A, axis=0)
return inter/(area_bbox+area_clusters-inter)
來計算一下 VOC2007 數據集的聚類結果，由於聚類結果會受到初始點的影響，所以每次運行的結果會不一樣，但是平均 IOU 應該是相近的：
root = 'data/VOCtrainval_06-Nov-2007/VOCdevkit/VOC2007/Annotations'
model = AnchorKmeans(root)
t0 = time()
clusters = model.get_cluster(9)
t1 = time()
print(f'耗時: {t1-t0} s')
print('聚類結果:\n', clusters*416)
print('平均 IOU:', model.average_iou(clusters))
特征圖分析
對於每個特征圖，我們給他分配了 3 個先驗框，如果數據集中有 C 個類，那么在特征圖的每個像素點處，我們會得到 3 個預測框，每個預測框預測一個 (4+1+C) 維的向量。假設我們使用 COCO 數據集，並將輸入神經網絡的圖片縮放為 416×416 的大小。由於我們的 COCO 數據集有 80 個類，所以在最小的特征圖的維度就應該是 13×13×(4+1+80)*3 即 13×13×255，為了方便理解和索引，我們將這個特征圖的維度 reshape 為 13×13×3×85。將 reshape 后的特征圖記為 y，那么 y[i, j, k] 的結果就會如下圖所示：
我們取到了一個 85 維的向量，上圖中這個向量的第1個元素 \(p_c\) 是 objectness score，代表這個預測框對包含物體的自信程度，我們將它重新記作 \(P(obj)\)。第 2 到 5 個元素 \((b_x, b_y, b_h, b_w)\) 表示這個預測框的中心點位置和寬高。后面 80 個元素 class probabilities 代表了在預測框含有物體的條件下，先驗框中是某個類別 \(c_i\) 的概率，所以這里的類別概率是一個條件概率 \(P(c_i|obj)\)，在我們使用 Yolov3 進行推理的時候，預測框上標注的概率是先驗概率，也就是 \(P(c_i)=P(obj)*P(c_i|obj)\)，對應到上圖就是那個 score（也用來 nms）。實際在代碼中我們會把邊界框的位置和寬高放到前 4 個元素，objectness score 放到第 5 個元素，后面的就是 class probabilities。
實際上我們不會直接預測 \((b_x, b_y, b_h, b_w)\)，而是預測偏差量 \((t_x,t_y,t_w,t_h)\) ，也就是說上圖中的第 2 到 5 個元素應該替換為偏差量。至於為什么要換成偏差量，原因應該是\((b_x, b_y, b_h, b_w)\) 的數值大小和 objectness score 以及 class probilities 差太多了，會給訓練帶來困難。這里先給出 \((t_x,t_y,t_w,t_h)\) 到 \((b_x, b_y, b_h, b_w)\) 的編碼公式，下面再詳細解讀：
\(c_x\) 和 \(c_y\) 是真實框的中心所處的那個單元格的坐標，\(t_x\) 和 \(t_y\) 是預測框和真實框的中心坐標偏差值。根據 Yolo 的思想：物體的中心落在哪個單元格，哪個單元格就應該負責預測這個物體。為了保證這個單元格正確預測，我們需要保證中心坐標偏差量不能超過 1，所以用 sigmoid 操作將 \(t_x\) 和 \(t_y\) 壓縮到 [0, 1] 區間內。\(p_w\) 和 \(p_h\) 是先驗框的寬高，\(t_w\) 和 \(t_h\) 代表了先驗框到真實框的縮放比。公式中多加了 exp 操作，應該是為了保證縮放比大於 0，不然在優化的時候會多一個 \(t_w>0\) 和 \(t_h>0\) 的約束，這時候 SGD 這種無約束求極值算法是用不了的。
訓練模型
損失函數
論文中沒有給出損失函數，但是可以根據各家代碼和博客中總結出 S×S 大小特征圖的損失函數來：
樣本划分
在 Yolov3 中規定了三種樣本：
- 正例，在真實框的中心點所在的區域內，與真實框交並比最大的那個先驗框就是正例，產生各種損失
- 負例，所有與真實框的交並比小於閾值的先驗框都是負例（不考慮交並比小於閾值的正例），只產生置信度損失 \(L_{obj}\)
- 忽視樣例，應該是為了平衡正負樣本，Yolov3 中將與真實框的交並比大於閾值的先驗框視為忽視樣例（不考慮交並比小於閾值的正例），不產生任何損失
划分完樣本之后，我們就知道上面的公式中 \(1_{i,j}^{obj}\) 和 \(1_{i,j}^{noobj}\) 的意思了，在正例的地方， \(1_{i,j}^{obj}\) 為 1，在反例的地方，\(1_{i,j}^{noobj}\) 為 1，否則為 0。
定位損失 \(L_{box}\)
只有正例可以產生定位損失，定位損失的公式中 \(\lambda_{coord}\) 是權重系數，\(2-w_i\times h_i\) 是一個懲罰因子，用來懲罰小預測框的損失。剩下的部分就是均方差損失。
置信度損失 \(L_{obj}\)
正例和負例都可以產生置信度損失，Yolov3 中直接將正例的置信度真值設置為 1，置信度損失使用交叉熵損失函數。
分類損失 \(L_{cls}\)
只有正例可以產生分類損失，分類損失直接使用了交叉熵損失函數。
代碼
樣本划分
def match(anchors: list, targets: List[Tensor], h: int, w: int, n_classes: int, overlap_thresh=0.5):
""" 匹配先驗框和邊界框真值
Parameters
----------
anchors: list of shape `(n_anchors, 2)`
根據特征圖的大小進行過縮放的先驗框
targets: List[Tensor]
標簽，每個元素的最后一個維度的第一個元素為類別，剩下四個為 `(cx, cy, w, h)`
h: int
特征圖的高度
w: int
特征圖的寬度
n_classes: int
類別數
overlap_thresh: float
IOU 閾值
Returns
-------
p_mask: Tensor of shape `(N, n_anchors, H, W)`
正例遮罩
n_mask: Tensor of shape `(N, n_anchors, H, W)`
反例遮罩
t: Tensor of shape `(N, n_anchors, H, W, n_classes+5)`
標簽
scale: Tensor of shape `(N, n_anchors, h, w)`
縮放值，用於懲罰小方框的定位
"""
N = len(targets)
n_anchors = len(anchors)
# 初始化返回值
p_mask = torch.zeros(N, n_anchors, h, w)
n_mask = torch.ones(N, n_anchors, h, w)
t = torch.zeros(N, n_anchors, h, w, n_classes+5)
scale = torch.zeros(N, n_anchors, h, w)
# 匹配先驗框和邊界框
anchors = np.hstack((np.zeros((n_anchors, 2)), np.array(anchors)))
for i in range(N):
target = targets[i] # shape:(n_objects, 5)
# 迭代每一個 ground truth box
for j in range(target.size(0)):
# 獲取標簽數據
cx, gw = target[j, [1, 3]]*w
cy, gh = target[j, [2, 4]]*h
# 獲取邊界框中心所處的單元格的坐標
gj, gi = int(cx), int(cy)
# 計算邊界框和先驗框的交並比
bbox = np.array([0, 0, gw, gh])
iou = jaccard_overlap_numpy(bbox, anchors)
# 標記出正例和反例
index = np.argmax(iou)
p_mask[i, index, gi, gj] = 1
# 正例除外，與 ground truth 的交並比都小於閾值則為負例
n_mask[i, index, gi, gj] = 0
n_mask[i, iou >= overlap_thresh, gi, gj] = 0
# 計算標簽值
t[i, index, gi, gj, 0] = cx-gj
t[i, index, gi, gj, 1] = cy-gi
t[i, index, gi, gj, 2] = math.log(gw/anchors[index, 2]+1e-16)
t[i, index, gi, gj, 3] = math.log(gh/anchors[index, 3]+1e-16)
t[i, index, gi, gj, 4] = 1
t[i, index, gi, gj, 5+int(target[j, 0])] = 1
# 縮放值，用於懲罰小方框的定位
scale[i, index, gi, gj] = 2-target[j, 3]*target[j, 4]
return p_mask, n_mask, t, scale
損失函數
# coding: utf-8
from typing import Tuple, List
import torch
from torch import Tensor, nn
from utils.box_utils import match
class YoloLoss(nn.Module):
""" 損失函數 """
def __init__(self, anchors: list, n_classes: int, image_size: int, overlap_thresh=0.5,
lambda_box=2.5, lambda_obj=1, lambda_noobj=0.5, lambda_cls=1):
"""
Parameters
----------
anchors: list of shape `(3, n_anchors, 2)`
先驗框列表
n_classes: int
類別數
image_size: int
輸入神經網絡的圖片大小
overlap_thresh: float
視為忽視樣例的 IOU 閾值
lambda_box, lambda_obj, lambda_noobj, lambda_cls: float
權重參數
"""
super().__init__()
self.anchors = anchors
self.n_classes = n_classes
self.image_size = image_size
self.lambda_box = lambda_box
self.lambda_obj = lambda_obj
self.lambda_noobj = lambda_noobj
self.lambda_cls = lambda_cls
self.overlap_thresh = overlap_thresh
self.mse_loss = nn.MSELoss(reduction='mean')
self.bce_loss = nn.BCELoss(reduction='mean')
def forward(self, preds: Tuple[Tensor], targets: List[Tensor]):
"""
Parameters
----------
preds: Tuple[Tensor]
Yolo 神經網絡輸出的各個特征圖，每個特征圖的維度為 `(N, (n_classes+5)*n_anchors, H, W)`
targets: List[Tensor]
標簽數據，每個標簽張量的維度為 `(N, n_objects, 5)`，最后一維的第一個元素為類別，剩下為邊界框 `(cx, cy, w, h)`
Returns
-------
loc_loss: Tensor
定位損失
conf_loss: Tensor
置信度損失
cls_loss: Tensor
分類損失
"""
loc_loss = 0
conf_loss = 0
cls_loss = 0
for anchors, pred in zip(self.anchors, preds):
N, _, img_h, img_w = pred.shape
n_anchors = len(anchors)
# 調整特征圖尺寸，方便索引
pred = pred.view(N, n_anchors, self.n_classes+5,
img_h, img_w).permute(0, 1, 3, 4, 2).contiguous()
# 獲取特征圖最后一個維度的每一部分
x = pred[..., 0].sigmoid()
y = pred[..., 1].sigmoid()
w = pred[..., 2]
h = pred[..., 3]
conf = pred[..., 4].sigmoid()
cls = pred[..., 5:].sigmoid()
# 匹配邊界框
step_h = self.image_size/img_h
step_w = self.image_size/img_w
anchors = [[i/step_w, j/step_h] for i, j in anchors]
p_mask, n_mask, t, scale = match(
anchors, targets, img_h, img_w, self.n_classes, self.overlap_thresh)
p_mask = p_mask.to(pred.device)
n_mask = n_mask.to(pred.device)
t = t.to(pred.device)
scale = scale.to(pred.device)
# 定位損失
x_loss = self.mse_loss(x*p_mask*scale, t[..., 0]*p_mask*scale)
y_loss = self.mse_loss(y*p_mask*scale, t[..., 1]*p_mask*scale)
w_loss = self.mse_loss(w*p_mask*scale, t[..., 2]*p_mask*scale)
h_loss = self.mse_loss(h*p_mask*scale, t[..., 3]*p_mask*scale)
loc_loss += (x_loss + y_loss + w_loss + h_loss)*self.lambda_box
# 置信度損失
conf_loss += self.bce_loss(conf*p_mask, p_mask)*self.lambda_obj + \
self.bce_loss(conf*n_mask, 0*n_mask)*self.lambda_noobj
# 分類損失
m = p_mask == 1
cls_loss += self.bce_loss(cls[m], t[..., 5:][m])*self.lambda_cls
return loc_loss, conf_loss, cls_loss
后記
Yolov3 的原理差不多就寫到這，實際上有各種各樣的實現方式，尤其是損失函數，這里就不一一介紹了。代碼放在了 GitHub，以上~