Method __init__

PATH/core/models/necks/simple_fpn.py:85–144

(self,
                 vis_token_dim,
                 mask_dim,
                 backbone,  # placeholder
                 bn_group,
                 pixel_decoder_cfg=None)

class SimpleFPN(nn.Module):
    def __init__(self,
                 vis_token_dim,
                 mask_dim,
                 backbone,  # placeholder
                 bn_group,
                 pixel_decoder_cfg=None):
        super(SimpleFPN, self).__init__()
        self.embed_dim = backbone.embed_dim
        self.mask_dim = mask_dim
        self.vis_token_dim = vis_token_dim
        self.pixel_decoder_cfg = pixel_decoder_cfg

        fpn1 = nn.Sequential(
            nn.ConvTranspose2d(self.embed_dim, self.embed_dim, kernel_size=2, stride=2),
            Norm2d(self.embed_dim),
            nn.GELU(),
            nn.ConvTranspose2d(self.embed_dim, self.embed_dim, kernel_size=2, stride=2),
            #  in compliance with decoder dim request
            Norm2d(self.embed_dim),
            nn.GELU(),
            nn.Conv2d(self.embed_dim, self.vis_token_dim, kernel_size=1, stride=1, padding=0),
            Norm2d(self.vis_token_dim),
        )

        fpn2 = nn.Sequential(
            nn.ConvTranspose2d(self.embed_dim, self.embed_dim, kernel_size=2, stride=2),
            #  in compliance with decoder dim request
            Norm2d(self.embed_dim),
            nn.GELU(),
            nn.Conv2d(self.embed_dim, self.vis_token_dim, kernel_size=1, stride=1, padding=0),
            Norm2d(self.vis_token_dim),
        )

        fpn3 = nn.Sequential(
            #  in compliance with decoder dim request
            nn.Conv2d(self.embed_dim, self.vis_token_dim, kernel_size=1, stride=1, padding=0),
            Norm2d(self.vis_token_dim),
        )

        fpn4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            #  in compliance with decoder dim request
            nn.Conv2d(self.embed_dim, self.vis_token_dim, kernel_size=1, stride=1, padding=0),
            Norm2d(self.vis_token_dim),
        )

        self.fpns = nn.ModuleList([fpn1, fpn2, fpn3, fpn4])

        if self.pixel_decoder_cfg is None:
            self.mask_features = nn.Conv2d(self.vis_token_dim, self.mask_dim, kernel_size=1, stride=1, padding=0)
            c2_xavier_fill(self.mask_features)
        else:
            input_shape = {name: ShapeSpec(channels=self.vis_token_dim, stride=[4, 8, 16, 32][i])
                           for i, name in enumerate(["fpn1", "fpn2", "fpn3", "fpn4"])}
            self.pixel_decoder = MSDeformAttnPixelDecoder(conv_dim=self.vis_token_dim,
                                                          input_shape=input_shape,
                                                          mask_dim=self.mask_dim,
                                                          **pixel_decoder_cfg)