ASFF(自适应空间特征融合）

两个不同尺度的特征进行融合时x1（2,128,640,640）x2（2,512,320,320），他们不应该是一视同仁的，一定有重要的不重要的，那就要找出权重来进行重新分配，融合成新的特征 x3=ax1’+bx2’

#--------------------------------------------#
#  上采样+通道变化(上采样倍率默认为2)
#--------------------------------------------#
class Upsample(nn.Module):
    """Applies convolution followed by upsampling."""
    def __init__(self, c1, c2, scale_factor=2):
        super().__init__()
        # self.cv1 = Conv(c1, c2, 1)
        # self.upsample = nn.Upsample(scale_factor=scale_factor, mode='nearest')  # or model='bilinear' non-deterministic
        if scale_factor == 2:
            self.cv1 = nn.ConvTranspose2d(c1, c2, 2, 2, 0, bias=True)  # 如果下采样率为2，就用Stride为2的2×2卷积来实现2次下采样
        elif scale_factor == 4:
            self.cv1 = nn.ConvTranspose2d(c1, c2, 4, 4, 0, bias=True)  # 如果下采样率为4，就用Stride为4的4×4卷积来实现4次下采样
 
    def forward(self, x):
        # return self.upsample(self.cv1(x))
        return self.cv1(x)

#--------------------------------------------#
#  自适应空间融合(ASFF)
#  eg:以 level1 为例子

#  如果特征x1,x2要进行ASFF 融合后到level1层
#  1.需要x2变成与x1通道，空间形状完全一样的x2'
#  2.之后将x1,x2'通道数分别进行降维到8， 
#  3.将压缩后的两个特征进行拼接后进行通道降维到2 ([b,2,h,w])
#  4.此时就已经取得了x1,x2'两个特征的不同权重
#  5.将权重拆分为两部分 分别乘x1,x2'，得到的特征再求和
#  6.最后经过一层卷积将融合的新特征的形状变得和x1一样
# （同理 level2 需要变得和x2形状一样）
#--------------------------------------------#
class ASFF(nn.Module):
    """ASFF2 module for YOLO AFPN head https://arxiv.org/abs/2306.15988"""
 
    def __init__(self, c1, level=0):
        super().__init__()
        #-----------------------------------#
        #  c1         :[c1,c2] 输入特征x1,x2 的通道
        #  level      :融合到哪一层中去
        #  inter_dim  :最终输出的通道数
        #  compress_c :中间被压缩的通道数
        #-----------------------------------#
        c1, c2 = c1[0], c1[1]
        self.level = level
        self.dim = c1, c2
        self.inter_dim = self.dim[self.level]
        compress_c = 8
       
        #------------------------------------#
        #  level1  x2变为和x1形状相同
        #  level2  x1变为和x2形状相同
        #------------------------------------#
        if level == 0:
            self.stride_level_1 = Upsample(c2, self.inter_dim)
        if level == 1:
            self.stride_level_0 = Conv(c1, self.inter_dim, 2, 2, 0)  # stride=2 下采样为2倍
 
  
        #------------------------------------#
        #  分别对两特征通道降维
        #------------------------------------#
        self.weight_level_0 = Conv(self.inter_dim, compress_c, 1, 1)
        self.weight_level_1 = Conv(self.inter_dim, compress_c, 1, 1)
        
        #------------------------------------#
        #  将降维后的特征合并并降维到2，得到每个特征的权重
        #------------------------------------#
        self.weights_levels = nn.Conv2d(compress_c * 2, 2, kernel_size=1, stride=1, padding=0)
        self.conv = Conv(self.inter_dim, self.inter_dim, 3, 1)
 
    def forward(self, x):
        #--------------------------------#
        #  获得两个特征
        #--------------------------------#
        x0, x1 = x[0], x[1]
 
        #--------------------------------#
        #  根据要融合到不同的level  x1/x2进行形状变化
        #--------------------------------#
        if self.level == 0:
            level_0_resized = x0
            level_1_resized = self.stride_level_1(x1)

        elif self.level == 1:
            level_0_resized = self.stride_level_0(x0)
            level_1_resized = x1
 
        #----------------------------------#
        #  通道降维
        #----------------------------------#
        level_0_weight_v = self.weight_level_0(level_0_resized)
        level_1_weight_v = self.weight_level_1(level_1_resized)
 
        #----------------------------------#
        #  特征合并并降维到2，得到不同特征的权重
        #----------------------------------#
        levels_weight_v = torch.cat((level_0_weight_v, level_1_weight_v), 1)
        levels_weight = self.weights_levels(levels_weight_v)

        #----------------------------------#
        #  权重归一化
        #----------------------------------#
        levels_weight = F.softmax(levels_weight, dim=1)
  
        #----------------------------------#
        #  将权重分配到不同的特征层
        #----------------------------------#
        fused_out_reduced = level_0_resized * levels_weight[:, 0:1] + level_1_resized * levels_weight[:, 1:2]
        return self.conv(fused_out_reduced)


model=ASFF2([128,512],0)
a=torch.Tensor(2,128,640,640)
b=torch.Tensor(2,512,320,320)
c=model([a,b])
print(c.shape)