yolox head
yolox head 網絡
概述
yolox-head 包含了3個分支
- 三個大分支輸入的是三種尺度特征圖，
- 自下而上分別對應前面提到的backbone輸出的dark3, dark4, dark5 。
- 尺度由大到小，堆疊成金字塔型。
代碼位置： yolox/models/yolo_head.py 下的 YOLOhead 類
主體結構
self.cls_convs = nn.ModuleList() # 兩個3x3的卷積
self.reg_convs = nn.ModuleList() # 兩個3x3的卷積
# pred
self.cls_preds = nn.ModuleList() # 一個1x1的卷積，把通道數變成類別數，比如coco 80類
self.reg_preds = nn.ModuleList() # 一個1x1的卷積，把通道數變成4通道，因為位置是xywh
self.obj_preds = nn.ModuleList() # 一個1x1的卷積，把通道數變成1通道，判斷有無目標
self.stems = nn.ModuleList() # 模前面的 BaseConv模塊
初始化head，__init__
# 3個不同尺度的輸出分支（對應dark3, dark4, dark5），期間用到的組件都是一樣的
for i in range(len(in_channels)):
self.stems.append(
# 開頭的 CBL 1x1 卷積 降維
BaseConv(
in_channels=int(in_channels[i] * width),
out_channels=int(256 * width),
ksize=1,
stride=1,
act=act,
)
)
# 兩個分支，分類分支和回歸分支
# 分類的卷積部分 開頭包含兩個卷積
self.cls_convs.append(
nn.Sequential(
*[
Conv(
in_channels=int(256 * width),
out_channels=int(256 * width),
ksize=3,
stride=1,
act=act,
),
Conv(
in_channels=int(256 * width),
out_channels=int(256 * width),
ksize=3,
stride=1,
act=act,
),
]
)
)
# 回歸的卷積部分 包含了兩層 卷積
self.reg_convs.append(
nn.Sequential(
*[
Conv(
in_channels=int(256 * width),
out_channels=int(256 * width),
ksize=3,
stride=1,
act=act,
),
Conv(
in_channels=int(256 * width),
out_channels=int(256 * width),
ksize=3,
stride=1,
act=act,
),
]
)
)
# 分類的預測部分 包含一層卷積，返回的是類別
self.cls_preds.append(
nn.Conv2d(
in_channels=int(256 * width),
out_channels=self.n_anchors * self.num_classes,
kernel_size=1,
stride=1,
padding=0,
)
)
# 回歸的預測部分 包含了一層卷積，返回的是四維度坐標
self.reg_preds.append(
nn.Conv2d(
in_channels=int(256 * width),
out_channels=4,
kernel_size=1,
stride=1,
padding=0,
)
)
# 目標的預測部分，返回的是目標是否存在
self.obj_preds.append(
nn.Conv2d(
in_channels=int(256 * width),
out_channels=self.n_anchors * 1,
kernel_size=1,
stride=1,
padding=0,
)
)
head模型的forward
for k, (cls_conv, reg_conv, stride_this_level, x) in enumerate(
zip(self.cls_convs, self.reg_convs, self.strides, xin)):
# CBL卷積
x = self.stems[k](x)
cls_x = x
reg_x = x
# 分類-卷積
cls_feat = cls_conv(cls_x)
# 分類-預測
cls_output = self.cls_preds[k](cls_feat)
# 回歸-卷積
reg_feat = reg_conv(reg_x)
# 回歸-預測
reg_output = self.reg_preds[k](reg_feat)
# 目標-預測
obj_output = self.obj_preds[k](reg_feat)
# 同一層合並
output = torch.cat(
[reg_output, obj_output.sigmoid(), cls_output.sigmoid()], 1)
# 不同層疊加
outputs.append(output)
結果輸出
concat + reshape + concat + transpose
- 三層特征經過
concat組件，是把分類和回歸結果按channel維度，即dim=1拼接； - 然后是
reshape，將特征圖展平成向量，(b,c,h,w) -> (b,c, h*w)，h*w即預測anchor個數； - 按dim=2，
concat不同尺度下的輸出，(b,c, h*w) -> (b,c, num_anchors)； - 然后是轉置操作，
(b,c, num_anchors) -> (b, num_anchors, c)。
# channel維度，將分類和回歸分支結果拼接。
output = torch.cat(
[reg_output, obj_output.sigmoid(), cls_output.sigmoid()], 1
)
# 1, reshape + concat + transpose, (b,c,h,w) -> (b,c,h*w) -> (b,c, ?) -> (b, ?, c)
outputs = torch.cat(
[x.flatten(start_dim=2) for x in outputs], dim=2
).permute(0, 2, 1)
轉置之后的輸出維度是(b, num_anchors, c)，其中每一行是一個預測的anchor信息。后面就是解碼，即將這些輸入翻譯成對應的預測框。
解碼
對網絡的輸出進行解碼，這里需要解碼信息是回歸的位置信息（分類信息不需要解碼），因為輸出的xywh是相對位置，簡單來說解碼過程就是(x+x_c, y+y_c, w, h) * stride，即預測的相對於網格左上角偏移的位置加上網格的位置，再乘以下采樣倍數，映射到原圖位置。解碼模塊的輸入是 (b, num_anchors, c)
def decode_outputs(self, outputs, dtype):
# outputs=(b, num_anchors, c)
grids = []
strides = []
# 計算每個尺度下所有網格的位置和對應的下采樣倍數
for (hsize, wsize), stride in zip(self.hw, self.strides):
# yv和xv分別存儲了每個網格的行和列。shape都是(hsize, wsize)
yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
# (hsize, wsize) -> (hsize, wsize, 2) -> (1, hsize*wsize, 2)
# 這樣每一行對應的是一個網絡的行列號。
grid = torch.stack((xv, yv), 2).view(1, -1, 2)
# 存儲每個尺度下所有網格的位置和對應的下采樣倍數
grids.append(grid)
shape = grid.shape[:2]
# (1, hsize*wsize, 1) 存儲放大倍數
strides.append(torch.full((*shape, 1), stride))
# 多個(1,hsize*wsize,2) -> (1,all_num_grids,2)，並轉換類型。主要是把所有不同尺度下的網格位置信息拼接起來。
grids = torch.cat(grids, dim=1).type(dtype)
# 同理。 多個(1,hsize*wsize,1) -> (1,all_num_grids,1)
strides = torch.cat(strides, dim=1).type(dtype)
# x,y位置偏移outputs[..., :2], shape=(1, all_num_grids, 2)
# grids所有網格的xy行列號, shape=(1, all_num_grids, 2)
# strides所有網格的下采樣倍數, shape=(1, all_num_grids, 1)
outputs[..., :2] = (outputs[..., :2] + grids) * strides
outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
return outputs
yolo-backbone 和 yolo-head 結合
地址：yolox/models/yolox.py 中的 YOLOX 類型
注意：
backbone輸出的是三個值- 在
head的forward第一個輸入參數也是三個值
class YOLOX(nn.Module):
"""
YOLOX model module. The module list is defined by create_yolov3_modules function.
The network returns loss values from three YOLO layers during training
and detection results during test.
"""
def __init__(self, backbone=None, head=None):
super().__init__()
if backbone is None:
backbone = YOLOPAFPN()
if head is None:
head = YOLOXHead(80)
self.backbone = backbone
self.head = head
def forward(self, x, targets=None):
# fpn output content features of [dark3, dark4, dark5]
# fpnout 輸出的是3個值
fpn_outs = self.backbone(x)
if self.training:
assert targets is not None
loss, iou_loss, conf_loss, cls_loss, l1_loss, num_fg = self.head(fpn_outs, targets, x)
outputs = {
"total_loss": loss,
"iou_loss": iou_loss,
"l1_loss": l1_loss,
"conf_loss": conf_loss,
"cls_loss": cls_loss,
"num_fg": num_fg,
}
else:
outputs = self.head(fpn_outs)
return outputs
參考博客：https://blog.csdn.net/jizhidexiaoming/article/details/119775002