FCOS加入ASFF的Attention机制

FCOS加入ASFF加权结构

这里是我对FCOS的结构加入了ASFF的Attention机制，将FCOS中fpn.py替换为下面代码即可。

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
import torch.nn.functional as F
from torch import nn


class ASFF(nn.Module):
    def __init__(self, level, rfb=False, vis=False):
        super(ASFF, self).__init__()
        self.level = level
        # self.dim = [512, 256, 256]
        self.dim = [256, 256, 256]
        self.inter_dim = self.dim[self.level]
        if level == 0:
            self.stride_level_1 = add_conv(256, self.inter_dim, 3, 2)
            self.stride_level_2 = add_conv(256, self.inter_dim, 3, 2)
            # self.expand = add_conv(self.inter_dim, 1024, 3, 1)
            self.expand = add_conv(self.inter_dim, 256, 3, 1)
        elif level == 1:
            # self.compress_level_0 = add_conv(512, self.inter_dim, 1, 1)
            self.compress_level_0 = add_conv(256, self.inter_dim, 1, 1)
            self.stride_level_2 = add_conv(256, self.inter_dim, 3, 2)
            # self.expand = add_conv(self.inter_dim, 512, 3, 1)
            self.expand = add_conv(self.inter_dim, 256, 3, 1)
        elif level == 2:
            # self.compress_level_0 = add_conv(512, self.inter_dim, 1, 1)
            self.compress_level_0 = add_conv(256, self.inter_dim, 1, 1)
            self.expand = add_conv(self.inter_dim, 256, 3, 1)

        # when adding rfb, we use half number of channels to save memory
        compress_c = 8 if rfb else 16

        self.weight_level_0 = add_conv(self.inter_dim, compress_c, 1, 1)
        self.weight_level_1 = add_conv(self.inter_dim, compress_c, 1, 1)
        self.weight_level_2 = add_conv(self.inter_dim, compress_c, 1, 1)

        self.weight_levels = nn.Conv2d(
            compress_c * 3, 3, kernel_size=1, stride=1, padding=0)
        self.vis = vis

    def forward(self, x_level_0, x_level_1, x_level_2):
        # import ipdb
        # ipdb.set_trace()
        if self.level == 0:
            level_0_resized = x_level_0
            level_1_resized =