Class Max_Former

event/max_former.py:13–97

__all__ = ['max_former']

class Max_Former(nn.Module):
    def __init__(self,
                 in_channels=2, num_classes=10,
                 embed_dims=[64, 128, 256], mlp_ratios=[4, 4, 4],
                 depths=[6, 8, 6], T = 4
                 ):
        super().__init__()

        self.num_classes = num_classes
        self.depths = depths
        self.T = T

        patch_embed1 = Embed_Max_plus( in_channels=in_channels,
                                       embed_dims=embed_dims // 2)

        stage1 = nn.ModuleList(  
            [Block_DWC3(
            dim=embed_dims // 2,  mlp_ratio=mlp_ratios)
            for j in range(1)]
            )


        patch_embed2 = Embed_Max( in_channels=embed_dims // 2,
                                    embed_dims=embed_dims)

        stage2 = nn.ModuleList([Block_SSA(
            dim=embed_dims,  mlp_ratio=mlp_ratios, num_heads = 16)
            for j in range(1)])

        setattr(self, f"patch_embed1", patch_embed1)
        setattr(self, f"patch_embed2", patch_embed2)
        setattr(self, f"stage1", stage1)
        setattr(self, f"stage2", stage2)
        
        self.head_lif = MultiStepLIFNode(tau=2.0, detach_reset=True)

        # classification head
        self.head = nn.Linear(embed_dims, num_classes) if num_classes > 0 else nn.Identity()
        self.apply(self._init_weights)

    @torch.jit.ignore
    def _get_pos_embed(self, pos_embed, patch_embed, H, W):
        if H * W == self.patch_embed1.num_patches:
            return pos_embed
        else:
            return F.interpolate(
                pos_embed.reshape(1, patch_embed.H, patch_embed.W, -1).permute(0, 3, 1, 2),
                size=(H, W), mode="bilinear").reshape(1, -1, H * W).permute(0, 2, 1)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)