cls_head_enhance_and_ECA
This commit is contained in:
Zhang Wei committed
1 parent
5dbde58605
commit
0bf2989264
3 files changed
+191
-5
No files matched your search
+127
-3
@@ -138,6 +138,24 @@ class ClsEnhanceBlock(nn.Module):
|
||||
def forward(self, x):
|
||||
return x + self.cv2(self.cv1(x))
|
||||
|
||||
# 新增
|
||||
class ECABlock(nn.Module):
|
||||
"""Efficient Channel Attention for G12."""
|
||||
def __init__(self, c1, k_size=3):
|
||||
super().__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
def forward(self, x):
|
||||
# x: [B, C, H, W]
|
||||
y = self.avg_pool(x) # [B, C, 1, 1]
|
||||
y = y.squeeze(-1).transpose(-1, -2) # [B, 1, C]
|
||||
y = self.conv(y) # [B, 1, C]
|
||||
y = self.sigmoid(y)
|
||||
y = y.transpose(-1, -2).unsqueeze(-1) # [B, C, 1, 1]
|
||||
return x * y.expand_as(x)
|
||||
|
||||
# 新增
|
||||
class DetectDecoupledLite(nn.Module):
|
||||
stride = None
|
||||
@@ -290,6 +308,87 @@ class DetectDecoupledLiteClsEnhance(nn.Module):
|
||||
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
|
||||
return grid, anchor_grid
|
||||
|
||||
#新增
|
||||
class DetectDecoupledLiteClsEnhanceECA(nn.Module):
|
||||
stride = None
|
||||
dynamic = False
|
||||
export = False
|
||||
|
||||
def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
|
||||
super().__init__()
|
||||
self.nc = nc
|
||||
self.no = nc + 5
|
||||
self.nl = len(anchors)
|
||||
self.na = len(anchors[0]) // 2
|
||||
self.grid = [torch.empty(0) for _ in range(self.nl)]
|
||||
self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]
|
||||
self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))
|
||||
self.inplace = inplace
|
||||
|
||||
head_c = 128
|
||||
|
||||
# stem
|
||||
self.stems = nn.ModuleList(Conv(c, head_c, 1, 1) for c in ch)
|
||||
|
||||
# reg + obj shared trunk
|
||||
self.regobj_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch)
|
||||
|
||||
# cls trunk
|
||||
self.cls_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch)
|
||||
|
||||
# G11 enhance
|
||||
self.cls_enhance = nn.ModuleList(ClsEnhanceBlock(head_c, expand=0.5) for _ in ch)
|
||||
|
||||
# G12 ECA attention
|
||||
self.cls_attn = nn.ModuleList(ECABlock(head_c, k_size=3) for _ in ch)
|
||||
|
||||
# predictors
|
||||
self.box_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * 4, 1) for _ in ch)
|
||||
self.obj_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * 1, 1) for _ in ch)
|
||||
self.cls_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * self.nc, 1) for _ in ch)
|
||||
|
||||
def forward(self, x):
|
||||
z = []
|
||||
for i in range(self.nl):
|
||||
xi = self.stems[i](x[i])
|
||||
|
||||
regobj_feat = self.regobj_convs[i](xi)
|
||||
|
||||
cls_feat = self.cls_convs[i](xi)
|
||||
cls_feat = self.cls_enhance[i](cls_feat)
|
||||
cls_feat = self.cls_attn[i](cls_feat)
|
||||
|
||||
box = self.box_preds[i](regobj_feat)
|
||||
obj = self.obj_preds[i](regobj_feat)
|
||||
cls = self.cls_preds[i](cls_feat)
|
||||
|
||||
out = torch.cat((box, obj, cls), 1)
|
||||
|
||||
bs, _, ny, nx = out.shape
|
||||
out = out.view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
||||
x[i] = out
|
||||
|
||||
if not self.training:
|
||||
if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
|
||||
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
|
||||
|
||||
xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
|
||||
xy = (xy * 2 + self.grid[i]) * self.stride[i]
|
||||
wh = (wh * 2) ** 2 * self.anchor_grid[i]
|
||||
y = torch.cat((xy, wh, conf), 4)
|
||||
z.append(y.view(bs, self.na * nx * ny, self.no))
|
||||
|
||||
return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
|
||||
|
||||
def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
|
||||
d = self.anchors[i].device
|
||||
t = self.anchors[i].dtype
|
||||
shape = 1, self.na, ny, nx, 2
|
||||
y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
|
||||
yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)
|
||||
grid = torch.stack((xv, yv), 2).expand(shape) - 0.5
|
||||
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
|
||||
return grid, anchor_grid
|
||||
|
||||
class Segment(Detect):
|
||||
"""YOLOv5 Segment head for segmentation models, extending Detect with mask and prototype layers."""
|
||||
@@ -373,7 +472,8 @@ class BaseModel(nn.Module):
|
||||
m = self.model[-1] # Detect()
|
||||
# if isinstance(m, (Detect, Segment)):
|
||||
# if isinstance(m, (Detect, DetectDecoupledLite, Segment)):
|
||||
if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)):
|
||||
# if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)):
|
||||
if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, DetectDecoupledLiteClsEnhanceECA, Segment)):
|
||||
m.stride = fn(m.stride)
|
||||
m.grid = list(map(fn, m.grid))
|
||||
if isinstance(m.anchor_grid, list):
|
||||
@@ -412,7 +512,8 @@ class DetectionModel(BaseModel):
|
||||
m = self.model[-1] # Detect()
|
||||
# if isinstance(m, (Detect, Segment)):
|
||||
# if isinstance(m, (Detect, DetectDecoupledLite, Segment)):
|
||||
if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)):
|
||||
# if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)):
|
||||
if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, DetectDecoupledLiteClsEnhanceECA, Segment)):
|
||||
|
||||
def _forward(x):
|
||||
"""Passes the input 'x' through the model and returns the processed output."""
|
||||
@@ -431,6 +532,8 @@ class DetectionModel(BaseModel):
|
||||
self._initialize_decoupled_biases()
|
||||
elif isinstance(m, DetectDecoupledLiteClsEnhance):
|
||||
self._initialize_decoupled_cls_enhance_biases()
|
||||
elif isinstance(m, DetectDecoupledLiteClsEnhanceECA):
|
||||
self._initialize_decoupled_cls_enhance_eca_biases()
|
||||
|
||||
# Init weights, biases
|
||||
initialize_weights(self)
|
||||
@@ -477,6 +580,26 @@ class DetectionModel(BaseModel):
|
||||
b.data += torch.log(cf / cf.sum())
|
||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
|
||||
# 新增
|
||||
def _initialize_decoupled_cls_enhance_eca_biases(self, cf=None):
|
||||
"""Initialize biases for DetectDecoupledLiteClsEnhanceECA."""
|
||||
m = self.model[-1]
|
||||
|
||||
# objectness bias
|
||||
for mi, s in zip(m.obj_preds, m.stride):
|
||||
b = mi.bias.view(m.na, -1)
|
||||
b.data += math.log(8 / (640 / s) ** 2)
|
||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
|
||||
# classification bias
|
||||
for mi in m.cls_preds:
|
||||
b = mi.bias.view(m.na, -1)
|
||||
if cf is None:
|
||||
b.data += math.log(0.6 / (m.nc - 0.99999))
|
||||
else:
|
||||
b.data += torch.log(cf / cf.sum())
|
||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
|
||||
def forward(self, x, augment=False, profile=False, visualize=False):
|
||||
"""Performs single-scale or augmented inference and may include profiling or visualization."""
|
||||
if augment:
|
||||
@@ -650,7 +773,8 @@ def parse_model(d, ch):
|
||||
# TODO: channel, gw, gd
|
||||
# elif m in {Detect, Segment}:
|
||||
# elif m in {Detect, DetectDecoupledLite, Segment}:
|
||||
elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment}:
|
||||
# elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment}:
|
||||
elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, DetectDecoupledLiteClsEnhanceECA, Segment}:
|
||||
args.append([ch[x] for x in f])
|
||||
if isinstance(args[1], int): # number of anchors
|
||||
args[1] = [list(range(args[1] * 2))] * len(f)
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
# Parameters
|
||||
nc: 100 # number of classes
|
||||
depth_multiple: 0.33 # model depth multiple
|
||||
width_multiple: 0.50 # layer channel multiple
|
||||
anchors:
|
||||
- [19, 27, 44, 40, 38, 94] # P3/8
|
||||
- [96, 68, 86, 152, 180, 137] # P4/16
|
||||
- [140, 301, 303, 264, 238, 542] # P5/32
|
||||
- [436, 615, 739, 380, 925, 792] # P6/64
|
||||
|
||||
# YOLOv5 v6.0 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[
|
||||
[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, C3, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 6, C3, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, C3, [512]],
|
||||
[-1, 1, Conv, [768, 3, 2]], # 7-P5/32
|
||||
[-1, 3, C3, [768]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
|
||||
[-1, 3, C3, [1024]],
|
||||
[-1, 1, SPPF, [1024, 5]], # 11
|
||||
]
|
||||
|
||||
# YOLOv5 v6.0 head
|
||||
head: [
|
||||
[-1, 1, Conv, [768, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, "nearest"]],
|
||||
[[-1, 8], 1, Concat, [1]], # cat backbone P5
|
||||
[-1, 3, C3, [768, False]], # 15
|
||||
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, "nearest"]],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 3, C3, [512, False]], # 19
|
||||
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, "nearest"]],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 3, C3, [256, False]], # 23 (P3/8-small)
|
||||
|
||||
[-1, 1, Conv, [256, 3, 2]],
|
||||
[[-1, 20], 1, Concat, [1]], # cat head P4
|
||||
[-1, 3, C3, [512, False]], # 26 (P4/16-medium)
|
||||
|
||||
[-1, 1, Conv, [512, 3, 2]],
|
||||
[[-1, 16], 1, Concat, [1]], # cat head P5
|
||||
[-1, 3, C3, [768, False]], # 29 (P5/32-large)
|
||||
|
||||
[-1, 1, Conv, [768, 3, 2]],
|
||||
[[-1, 12], 1, Concat, [1]], # cat head P6
|
||||
[-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
|
||||
|
||||
# [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
|
||||
[[23, 26, 29, 32], 1, DetectDecoupledLiteClsEnhanceECA, [nc, anchors]], # Detect(P3, P4, P5, P6)
|
||||
]
|
||||
Reference in new issue
Block a user