如何在 YOLOv7 中添加 GAM 注意力机制:代码实现
以下是在 YOLOv7 中添加 GAM 注意力机制的代码实现:
- 首先,需要在 model.py 中添加 GAM 注意力机制的类:
class GAM(nn.Module):
def __init__(self, in_channels, out_channels, reduction=16):
super(GAM, self).__init__()
self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
self.conv1 = nn.Conv2d(in_channels, out_channels // reduction, kernel_size=1, bias=False)
self.conv2 = nn.Conv2d(out_channels // reduction, out_channels, kernel_size=1, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
b, c, _, _ = x.size()
y = self.global_avg_pool(x)
y = self.conv1(y)
y = F.relu(y, inplace=True)
y = self.conv2(y)
y = self.sigmoid(y)
return x * y.expand_as(x)
- 在 CSPDarknet 中添加 GAM 注意力机制:
class CSPBlock(nn.Module):
def __init__(self, in_channels, out_channels, num_blocks, use_gam=True):
super(CSPBlock, self).__init__()
self.split = nn.Sequential(
Conv(in_channels, out_channels, kernel_size=3, stride=2),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
*[Residual(out_channels, use_gam=use_gam) for _ in range(num_blocks)]
)
self.downsample = nn.Sequential(
Conv(out_channels, out_channels, kernel_size=1, stride=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
Conv(out_channels, out_channels, kernel_size=3, stride=2),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
*[Residual(out_channels, use_gam=use_gam) for _ in range(num_blocks)],
)
self.merge = nn.Sequential(
Conv(2 * out_channels, out_channels, kernel_size=1, stride=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
*[Residual(out_channels, use_gam=use_gam) for _ in range(num_blocks)]
)
def forward(self, x):
x1 = self.split(x)
x2 = self.downsample(x1)
x = torch.cat([x1, x2], dim=1)
x = self.merge(x)
return x
在 Residual 模块中添加 GAM 注意力机制:
class Residual(nn.Module):
def __init__(self, channels, use_gam=True):
super(Residual, self).__init__()
self.use_gam = use_gam
self.conv1 = Conv(channels, channels, kernel_size=1, stride=1)
self.bn1 = nn.BatchNorm2d(channels)
self.conv2 = Conv(channels, channels, kernel_size=3, stride=1)
self.bn2 = nn.BatchNorm2d(channels)
self.conv3 = Conv(channels, channels, kernel_size=1, stride=1)
self.bn3 = nn.BatchNorm2d(channels)
if self.use_gam:
self.gam = GAM(channels, channels)
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = F.relu(out, inplace=True)
out = self.conv2(out)
out = self.bn2(out)
out = F.relu(out, inplace=True)
out = self.conv3(out)
out = self.bn3(out)
if self.use_gam:
out = self.gam(out)
out += residual
out = F.relu(out, inplace=True)
return out
- 在 YOLOv7 中使用 CSPDarknet with GAM 注意力机制:
class YOLOv7(nn.Module):
def __init__(self, num_classes=80):
super(YOLOv7, self).__init__()
self.backbone = CSPDarknet(use_gam=True)
self.neck = SPP(use_gam=True)
self.head = YOLOHead(in_channels=[128, 256, 512], num_anchors=3, num_classes=num_classes)
def forward(self, x):
x = self.backbone(x)
x = self.neck(x)
x = self.head(x)
return x
这样,就在 YOLOv7 中添加了 GAM 注意力机制。
原文地址: https://www.cveoy.top/t/topic/m4mC 著作权归作者所有。请勿转载和采集!