在深度学习中,注意力机制(Attention Mechanism) 是让模型学会"关注重点"的方法。正如人类 在看图时会自动聚焦于主体(如猫、车、人脸),而忽略背景,模型也希望学会同样的能力。
python
class SimpleCNN(nn.Module):
def __init__(self, num_classes=10):
super(SimpleCNN, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
SEBlock(64) # 加入通道注意力模块
)
self.layer2 = nn.Sequential(
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
SEBlock(128)
)
self.fc = nn.Linear(128, num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = torch.mean(out, dim=[2, 3]) # Global AvgPool
out = self.fc(out)
return out
model = SimpleCNN()
print(model)
python
if __name__ == "__main__":
# 1. 加载图片
img_tensor, img = load_image()
# 2. 初始化两个CNN(加SE和不加SE)
cnn_no_se = SimpleCNN()
cnn_with_se = CNNWithSE()
cnn_no_se.eval()
cnn_with_se.eval()
# 3. 提取特征
feat1_no_se, feat2_no_se = cnn_no_se(img_tensor)
feat1_with_se, feat2_with_se = cnn_with_se(img_tensor)
# 4. 可视化对比(第一层特征图)
visualize_feat(feat1_no_se, "不加SE的第一层特征图")
visualize_feat(feat1_with_se, "加SE的第一层特征图")