change yolo.py
This commit is contained in:
Zhang Wei committed
1 parent
9b008eaa99
commit
b586eee607
1 file changed
+99
-10
+99
-10
@@ -68,6 +68,30 @@ try:
|
||||
except ImportError:
|
||||
thop = None
|
||||
|
||||
# 新增
|
||||
class ECAAttention(nn.Module):
|
||||
"""Efficient Channel Attention for lightweight channel re-weighting."""
|
||||
def __init__(self, k_size=3):
|
||||
super().__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.conv = nn.Conv1d(
|
||||
in_channels=1,
|
||||
out_channels=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 = y.transpose(-1, -2).unsqueeze(-1) # [B, C, 1, 1]
|
||||
y = self.sigmoid(y)
|
||||
return x * y.expand_as(x)
|
||||
|
||||
|
||||
class Detect(nn.Module):
|
||||
"""YOLOv5 Detect head for processing input tensors and generating detection outputs in object detection models."""
|
||||
@@ -126,6 +150,60 @@ class Detect(nn.Module):
|
||||
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
|
||||
return grid, anchor_grid
|
||||
|
||||
# 新增
|
||||
class ECADetect(nn.Module):
|
||||
"""YOLOv5 Detect head with ECA attention before prediction conv."""
|
||||
stride = None
|
||||
dynamic = False
|
||||
export = False
|
||||
|
||||
def __init__(self, nc=80, anchors=(), ch=(), inplace=True, eca_k=3):
|
||||
super().__init__()
|
||||
self.nc = nc # number of classes
|
||||
self.no = nc + 5 # number of outputs per anchor
|
||||
self.nl = len(anchors) # number of detection layers
|
||||
self.na = len(anchors[0]) // 2 # number of anchors
|
||||
self.grid = [torch.empty(0) for _ in range(self.nl)] # init grid
|
||||
self.anchor_grid = [torch.empty(0) for _ in range(self.nl)] # init anchor grid
|
||||
self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))
|
||||
self.inplace = inplace
|
||||
|
||||
# one ECA module for each detection scale
|
||||
self.eca = nn.ModuleList(ECAAttention(k_size=eca_k) for _ in ch)
|
||||
|
||||
# same prediction conv as original YOLOv5 Detect
|
||||
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)
|
||||
|
||||
def forward(self, x):
|
||||
z = [] # inference output
|
||||
for i in range(self.nl):
|
||||
x[i] = self.eca[i](x[i]) # ECA enhancement
|
||||
x[i] = self.m[i](x[i]) # prediction conv
|
||||
|
||||
bs, _, ny, nx = x[i].shape
|
||||
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
||||
|
||||
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."""
|
||||
@@ -207,7 +285,8 @@ class BaseModel(nn.Module):
|
||||
"""
|
||||
self = super()._apply(fn)
|
||||
m = self.model[-1] # Detect()
|
||||
if isinstance(m, (Detect, Segment)):
|
||||
# if isinstance(m, (Detect, Segment)):
|
||||
if isinstance(m, (Detect, ECADetect, Segment)):
|
||||
m.stride = fn(m.stride)
|
||||
m.grid = list(map(fn, m.grid))
|
||||
if isinstance(m.anchor_grid, list):
|
||||
@@ -244,7 +323,8 @@ class DetectionModel(BaseModel):
|
||||
|
||||
# Build strides, anchors
|
||||
m = self.model[-1] # Detect()
|
||||
if isinstance(m, (Detect, Segment)):
|
||||
# if isinstance(m, (Detect, Segment)):
|
||||
if isinstance(m, (Detect, ECADetect, Segment)):
|
||||
|
||||
def _forward(x):
|
||||
"""Passes the input 'x' through the model and returns the processed output."""
|
||||
@@ -322,13 +402,21 @@ class DetectionModel(BaseModel):
|
||||
"""
|
||||
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
|
||||
m = self.model[-1] # Detect() module
|
||||
for mi, s in zip(m.m, m.stride): # from
|
||||
b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
|
||||
b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
|
||||
b.data[:, 5 : 5 + m.nc] += (
|
||||
math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
|
||||
) # cls
|
||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
# for mi, s in zip(m.m, m.stride): # from
|
||||
# b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
|
||||
# b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
|
||||
# b.data[:, 5 : 5 + m.nc] += (
|
||||
# math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
|
||||
# ) # cls
|
||||
# mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
if isinstance(m, (Detect, ECADetect)):
|
||||
for mi, s in zip(m.m, m.stride):
|
||||
b = mi.bias.view(m.na, -1)
|
||||
b.data[:, 4] += math.log(8 / (640 / s) ** 2)
|
||||
b.data[:, 5:5 + m.nc] += (
|
||||
math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
|
||||
)
|
||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
|
||||
|
||||
Model = DetectionModel # retain YOLOv5 'Model' class for backwards compatibility
|
||||
@@ -434,7 +522,8 @@ def parse_model(d, ch):
|
||||
elif m is Concat:
|
||||
c2 = sum(ch[x] for x in f)
|
||||
# TODO: channel, gw, gd
|
||||
elif m in {Detect, Segment}:
|
||||
# elif m in {Detect, Segment}:
|
||||
elif m in {Detect, ECADetect, 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)
|
||||
|
||||
Reference in new issue
Block a user