计算机视觉之 SE 注意力模块

计算机视觉之 SE 注意力模块

一、简介

SEBlock 是一个自定义的神经网络模块，主要用于实现 Squeeze-and-Excitation（SE）注意力机制。SE 注意力机制通过全局平均池化和全连接层来重新校准通道的权重，从而增强模型的表达能力。

原论文：《Squeeze-and-Excitation Networks》

二、语法和参数

语法

class SEBlock(nn.Module):def __init__(self, in_channels, reduction=16):...def forward(self, x):...

参数

• in_channels：输入特征的通道数。
• reduction：通道缩减比例，默认为 16。

三、实例

3.1 初始化和前向传播

• 代码

import torch
import torch.nn as nnclass SEBlock(nn.Module):def __init__(self, in_channels, reduction=16):super(SEBlock, self).__init__()reduced_channels = max(in_channels // reduction, 1)self.global_avg_pool = nn.AdaptiveAvgPool2d(1)self.fc = nn.Sequential(nn.Linear(in_channels, reduced_channels, bias=False),nn.ReLU(inplace=True),nn.Linear(reduced_channels, in_channels, bias=False),nn.Sigmoid())def forward(self, x):batch_size, channels, _, _ = x.size()# Squeezey = self.global_avg_pool(x).view(batch_size, channels)# Excitationy = self.fc(y).view(batch_size, channels, 1, 1)# Scalereturn x * y.expand_as(x)

• 输出

加权图像输出

3.2 应用在示例数据上

• 代码

import torch# 创建示例输入数据
input_tensor = torch.randn(1, 64, 32, 32)  # (batch_size, in_channels, height, width)# 初始化 SEBlock 模块
se_block = SEBlock(in_channels=64, reduction=16)# 前向传播
output_tensor = se_block(input_tensor)
print(output_tensor.shape)

• 输出

torch.Size([1, 64, 32, 32])

四、注意事项

1. SEBlock 模块通过全局平均池化和全连接层来重新校准通道的权重，从而增强模型的表达能力。
2. 在使用 SEBlock 时，确保输入特征的通道数和缩减比例设置合理，以避免计算开销过大。
3. 该模块主要用于图像数据处理，适用于各种计算机视觉任务，如图像分类、目标检测等。

​