ResNet改进(51)：基于轴向注意力机制的改进ResNet模型

1.创新点解析

在计算机视觉领域，卷积神经网络(CNN)一直是图像识别任务的主流架构。

近年来，注意力机制的引入为CNN带来了新的改进方向。

本文将详细解析一个结合了轴向注意力(Axial Attention)机制的改进ResNet模型，展示如何通过注意力机制增强传统CNN的性能。

模型概述

我们构建的CustomResNet基于经典的ResNet34架构，但进行了两个重要改进：

1. 修改了初始卷积层，使用更小的卷积核(3x3)和步长1，保留更多空间信息

2. 在残差块之间插入了轴向注意力模块，增强模型对长距离依赖关系的捕捉能力

轴向注意力模块

轴向注意力是自注意力机制的一种高效变体，它分别沿着高度和宽度两个轴向计算注意力，显著降低了计算复杂度。

class AxialAttention(nn.Module):
    def __init__(self, in_channels, heads=8):
        super(AxialAttention, self).__init__()
        self.heads = heads
        self.scale = (in_channels