Day.46
通道注意力机制:
class ChannelAttention(nn.Module):
def __init__(self, in_channels, reduction_ratio=16):
in_channels: 输入特征图的通道数
reduction_ratio: 降维比例,用于减少参数量
super(ChannelAttention, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(in_channels, in_channels // reduction_ratio, bias=False),
nn.ReLU(inplace=True),
nn.Linear(in_channels // reduction_ratio, in_channels, bias=False),
nn.Sigmoid()
)
def forward(self, x):
"""
参数:
x: 输入特征图,形状为 [batch_size, channels, height, width]
返回:
加权后的特征图,形状不变
"""
batch_size, channels, height, width = x.size()
avg_pool_output = self.avg_pool(x)
avg_pool_output = avg_pool_output.view(batch_size, channels)
channel_weights = self.fc(avg_pool_output)
channel_weights = channel_weights.view(batch_size, channels, 1, 1)
return x * channel_weights
模型定义:
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
# ---------------------- 第一个卷积块 ----------------------
self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.relu1 = nn.ReLU()
self.ca1 = ChannelAttention(in_channels=32, reduction_ratio=16)
self.pool1 = nn.MaxPool2d(2, 2)
# ---------------------- 第二个卷积块 ----------------------
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.bn2 = nn.BatchNorm2d(64)
self.relu2 = nn.ReLU()
self.ca2 = ChannelAttention(in_channels=64, reduction_ratio=16)
self.pool2 = nn.MaxPool2d(2)
# ---------------------- 第三个卷积块 ----------------------
self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
self.bn3 = nn.BatchNorm2d(128)
self.relu3 = nn.ReLU()
self.ca3 = ChannelAttention(in_channels=128, reduction_ratio=16)
self.pool3 = nn.MaxPool2d(2)
# ---------------------- 全连接层(分类器) ----------------------
self.fc1 = nn.Linear(128 * 4 * 4, 512)
self.dropout = nn.Dropout(p=0.5)
self.fc2 = nn.Linear(512, 10)
def forward(self, x):
# ---------- 卷积块1处理 ----------
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.ca1(x)
x = self.pool1(x)
# ---------- 卷积块2处理 ----------
x = self.conv2(x)
x = self.bn2(x)
x = self.relu2(x)
x = self.ca2(x)
x = self.pool2(x)
# ---------- 卷积块3处理 ----------
x = self.conv3(x)
x = self.bn3(x)
x = self.relu3(x)
x = self.ca3(x)
x = self.pool3(x)
# ---------- 展平与全连接层 ----------
x = x.view(-1, 128 * 4 * 4)
x = self.fc1(x)
x = self.relu3(x)
x = self.dropout(x)
x = self.fc2(x)
return x
# 重新初始化模型,包含通道注意力模块
model = CNN()
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler.step()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='min',
patience=3,
factor=0.5 @浙大疏锦行
)