try modify swin

models.py

@@ -10,8 +10,6 @@ from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from timm.models.layers import DropPath, trunc_normal_
 from timm.models.registry import register_model
 from timm.layers.helpers import to_2tuple
-from typing import *
-import math
 
 
 class ConvE(torch.nn.Module):
@@ -437,6 +435,50 @@ class TuckER(torch.nn.Module):
 
         return pred
 
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(2 * dim)
+
+    def forward(self, x):
+        """
+        x: B, C, H, W
+        """
+        B, C, H, W = x.shape
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.reduction(x)
+        x = self.norm(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        flops += H * W * self.dim // 2
+        return flops
 
 class FouriER(torch.nn.Module):
     def __init__(self, params, hid_drop = None, embed_dim = None):
@@ -490,9 +532,10 @@ class FouriER(torch.nn.Module):
         self.patch_embed = PatchEmbed(in_chans=channels, patch_size=self.p.patch_size, 
                                       embed_dim=self.p.embed_dim, stride=4, padding=2)
         network = []
-        layers = [4, 4, 12, 4]
-        embed_dims = [self.p.embed_dim, 128, 320, 128]
-        mlp_ratios = [4, 4, 4, 4]
+        layers = [2, 2, 6, 2]
+        embed_dims = [self.p.embed_dim, 320, 256, 128]
+        mlp_ratios = [4, 4, 8, 12]
+        num_heads = [4, 4, 4, 4]
         downsamples = [True, True, True, True]
         pool_size=3
         act_layer=nn.GELU
@@ -504,6 +547,7 @@ class FouriER(torch.nn.Module):
         down_patch_size=3
         down_stride=2
         down_pad=1
+        window_size = 4
         num_classes=self.p.embed_dim
         for i in range(len(layers)):
             stage = basic_blocks(embed_dims[i], i, layers, 
@@ -512,7 +556,9 @@ class FouriER(torch.nn.Module):
                                  drop_rate=drop_rate, 
                                  drop_path_rate=drop_path_rate,
                                  use_layer_scale=use_layer_scale, 
-                                 layer_scale_init_value=layer_scale_init_value)
+                                 layer_scale_init_value=layer_scale_init_value,
+                                 num_heads=num_heads[i], input_resolution=(image_h // (2**i), image_w // (2**i)),
+                                 window_size=window_size, shift_size=0)
             network.append(stage)
             if i >= len(layers) - 1:
                 break
@@ -524,26 +570,11 @@ class FouriER(torch.nn.Module):
                         padding=down_pad, 
                         in_chans=embed_dims[i], embed_dim=embed_dims[i+1]
                         )
+                    # PatchMerging(dim=embed_dims[i+1])
                     )
 
         self.network = nn.ModuleList(network)
         self.norm = norm_layer(embed_dims[-1])
-        self.graph_type = 'Spatial'
-        N = (image_h // patch_size)**2
-        if self.graph_type in ["Spatial", "Mixed"]:
-            # Create a range tensor of node indices
-            indices = torch.arange(N)
-            # Reshape the indices tensor to create a grid of row and column indices
-            row_indices = indices.view(-1, 1).expand(-1, N)
-            col_indices = indices.view(1, -1).expand(N, -1)
-            # Compute the adjacency matrix
-            row1, col1 = row_indices // int(math.sqrt(N)), row_indices % int(math.sqrt(N))
-            row2, col2 = col_indices // int(math.sqrt(N)), col_indices % int(math.sqrt(N))
-            graph = ((abs(row1 - row2) <= 1).float() * (abs(col1 - col2) <= 1).float())
-            graph = graph - torch.eye(N)
-            self.spatial_graph = graph.cuda() # comment .to("cuda") if the environment is cpu
-        self.class_token = False
-        self.token_scale = False
         self.head = nn.Linear(
                 embed_dims[-1], num_classes) if num_classes > 0 \
                 else nn.Identity()
@@ -561,45 +592,8 @@ class FouriER(torch.nn.Module):
 
     def forward_tokens(self, x):
         outs = []
-        B, C, H, W = x.shape
-        N = H*W
-        if self.graph_type in ["Semantic", "Mixed"]:
-            # Generate the semantic graph w.r.t. the cosine similarity between tokens
-            # Compute cosine similarity
-            if self.class_token:
-                x_normed = x[:, 1:] / x[:, 1:].norm(dim=-1, keepdim=True)
-            else:
-                x_normed = x / x.norm(dim=-1, keepdim=True)
-            x_cossim = x_normed @ x_normed.transpose(-1, -2)
-            threshold = torch.kthvalue(x_cossim, N-1-self.num_neighbours, dim=-1, keepdim=True)[0] # B,H,1,1 
-            semantic_graph = torch.where(x_cossim>=threshold, 1.0, 0.0)
-            if self.class_token:
-                semantic_graph = semantic_graph - torch.eye(N-1, device=semantic_graph.device).unsqueeze(0)
-            else:
-                semantic_graph = semantic_graph - torch.eye(N, device=semantic_graph.device).unsqueeze(0)
-        
-        if self.graph_type == "None":
-            graph = None
-        else:
-            if self.graph_type == "Spatial":
-                graph = self.spatial_graph.unsqueeze(0).expand(B,-1,-1)#.to(x.device)
-            elif self.graph_type == "Semantic":
-                graph = semantic_graph
-            elif self.graph_type == "Mixed":
-                # Integrate the spatial graph and semantic graph
-                spatial_graph = self.spatial_graph.unsqueeze(0).expand(B,-1,-1).to(x.device)
-                graph = torch.bitwise_or(semantic_graph.int(), spatial_graph.int()).float()
-            
-            # Symmetrically normalize the graph
-            degree = graph.sum(-1) # B, N
-            degree = torch.diag_embed(degree**(-1/2))
-            graph = degree @ graph @ degree
-
         for idx, block in enumerate(self.network):
-            try:
-                x = block(x, graph)
-            except:
-                x = block(x)
+            x = block(x)
         # output only the features of last layer for image classification
         return x
 
@@ -742,7 +736,7 @@ def basic_blocks(dim, index, layers,
                  pool_size=3, mlp_ratio=4., 
                  act_layer=nn.GELU, norm_layer=GroupNorm, 
                  drop_rate=.0, drop_path_rate=0., 
-                 use_layer_scale=True, layer_scale_init_value=1e-5):
+                 use_layer_scale=True, layer_scale_init_value=1e-5, num_heads = 4, input_resolution = None, window_size = 4, shift_size = 2):
     """
     generate PoolFormer blocks for a stage
     return: PoolFormer blocks 
@@ -757,8 +751,10 @@ def basic_blocks(dim, index, layers,
             drop=drop_rate, drop_path=block_dpr, 
             use_layer_scale=use_layer_scale, 
             layer_scale_init_value=layer_scale_init_value, 
+            num_heads=num_heads, input_resolution = input_resolution,
+            window_size=window_size, shift_size=shift_size
             ))
-    blocks = SeqModel(*blocks)
+    blocks = nn.Sequential(*blocks)
 
     return blocks
 
@@ -876,9 +872,12 @@ class WindowAttention(nn.Module):
         attn = attn + relative_position_bias.unsqueeze(0)
 
         if mask is not None:
-            nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
-            attn = attn.view(-1, self.num_heads, N, N)
+            try:
+                nW = mask.shape[0]
+                attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+                attn = attn.view(-1, self.num_heads, N, N)
+            except:
+                pass
             attn = self.softmax(attn)
         else:
             attn = self.softmax(attn)
@@ -923,279 +922,6 @@ def window_reverse(windows, window_size, H, W):
     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, -1, H, W)
     return x
 
-class SeqModel(nn.Sequential):
-	def forward(self, *inputs):
-		for module in self._modules.values():
-			if type(inputs) == tuple:
-				inputs = module(*inputs)
-			else:
-				inputs = module(inputs)
-		return inputs
-
-def propagate(x: torch.Tensor, weight: torch.Tensor, 
-              index_kept: torch.Tensor, index_prop: torch.Tensor, 
-              standard: str = "None", alpha: Optional[float] = 0, 
-              token_scales: Optional[torch.Tensor] = None,
-              cls_token=True):
-    """
-    Propagate tokens based on the selection results.
-    ================================================
-    Args:
-        - x: Tensor([B, N, C]): the feature map of N tokens, including the [CLS] token.
-        
-        - weight: Tensor([B, N-1, N-1]): the weight of each token propagated to the other tokens, 
-                                         excluding the [CLS] token. weight could be a pre-defined 
-                                         graph of the current feature map (by default) or the
-                                         attention map (need to manually modify the Block Module).
-                                         
-        - index_kept: Tensor([B, N-1-num_prop]): the index of kept image tokens in the feature map X
-        
-        - index_prop: Tensor([B, num_prop]): the index of propagated image tokens in the feature map X
-        
-        - standard: str: the method applied to propagate the tokens, including "None", "Mean" and 
-                         "GraphProp"
-        
-        - alpha: float: the coefficient of propagated features
-        
-        - token_scales: Tensor([B, N]): the scale of tokens, including the [CLS] token. token_scales
-                                        is None by default. If it is not None, then token_scales 
-                                        represents the scales of each token and should sum up to N.
-        
-    Return:
-        - x: Tensor([B, N-1-num_prop, C]): the feature map after propagation
-        
-        - weight: Tensor([B, N-1-num_prop, N-1-num_prop]): the graph of feature map after propagation
-        
-        - token_scales: Tensor([B, N-1-num_prop]): the scale of tokens after propagation
-    """
-    
-    B, N, C = x.shape
-    
-    # Step 1: divide tokens
-    if cls_token:
-        x_cls = x[:, 0:1] # B, 1, C
-    x_kept = x.gather(dim=1, index=index_kept.unsqueeze(-1).expand(-1,-1,C)) # B, N-1-num_prop, C
-    x_prop = x.gather(dim=1, index=index_prop.unsqueeze(-1).expand(-1,-1,C)) # B, num_prop, C
-    
-    # Step 2: divide token_scales if it is not None
-    if token_scales is not None:
-        if cls_token:
-            token_scales_cls = token_scales[:, 0:1] # B, 1
-        token_scales_kept = token_scales.gather(dim=1, index=index_kept) # B, N-1-num_prop
-        token_scales_prop = token_scales.gather(dim=1, index=index_prop) # B, num_prop
-    
-    # Step 3: propagate tokens
-    if standard == "None":
-        """
-        No further propagation
-        """
-        pass
-        
-    elif standard == "Mean":
-        """
-        Calculate the mean of all the propagated tokens,
-        and concatenate the result token back to kept tokens.
-        """
-        # naive average
-        x_prop = x_prop.mean(1, keepdim=True) # B, 1, C
-        # Concatenate the average token 
-        x_kept = torch.cat((x_kept, x_prop), dim=1) # B, N-num_prop, C
-            
-    elif standard == "GraphProp":
-        """
-        Propagate all the propagated token to kept token
-        with respect to the weights and token scales.
-        """
-        assert weight is not None, "The graph weight is needed for graph propagation"
-        
-        # Step 3.1: divide propagation weights.
-        if cls_token:
-            index_kept = index_kept - 1 # since weights do not include the [CLS] token
-            index_prop = index_prop - 1 # since weights do not include the [CLS] token
-            weight = weight.gather(dim=1, index=index_kept.unsqueeze(-1).expand(-1,-1,N-1)) # B, N-1-num_prop, N-1
-            weight_prop = weight.gather(dim=2, index=index_prop.unsqueeze(1).expand(-1,weight.shape[1],-1)) # B, N-1-num_prop, num_prop
-            weight = weight.gather(dim=2, index=index_kept.unsqueeze(1).expand(-1,weight.shape[1],-1)) # B, N-1-num_prop, N-1-num_prop
-        else:
-            weight = weight.gather(dim=1, index=index_kept.unsqueeze(-1).expand(-1,-1,N)) # B, N-1-num_prop, N-1
-            weight_prop = weight.gather(dim=2, index=index_prop.unsqueeze(1).expand(-1,weight.shape[1],-1)) # B, N-1-num_prop, num_prop
-            weight = weight.gather(dim=2, index=index_kept.unsqueeze(1).expand(-1,weight.shape[1],-1)) # B, N-1-num_prop, N-1-num_prop
-        
-        # Step 3.2: generate the broadcast message and propagate the message to corresponding kept tokens
-        # Simple implementation
-        x_prop = weight_prop @ x_prop # B, N-1-num_prop, C
-        x_kept = x_kept + alpha * x_prop # B, N-1-num_prop, C
-        
-        """ scatter_reduce implementation for batched inputs
-        # Get the non-zero values
-        non_zero_indices = torch.nonzero(weight_prop, as_tuple=True)
-        non_zero_values = weight_prop[non_zero_indices]
-        
-        # Sparse multiplication
-        batch_indices, row_indices, col_indices = non_zero_indices
-        sparse_matmul = alpha * non_zero_values[:, None] * x_prop[batch_indices, col_indices, :]
-        reduce_indices = batch_indices * x_kept.shape[1] + row_indices
-        
-        x_kept = x_kept.reshape(-1, C).scatter_reduce(dim=0, 
-                                                      index=reduce_indices[:, None], 
-                                                      src=sparse_matmul, 
-                                                      reduce="sum",
-                                                      include_self=True)
-        x_kept = x_kept.reshape(B, -1, C)
-        """
-        
-        # Step 3.3: calculate the scale of each token if token_scales is not None
-        if token_scales is not None:
-            if cls_token:
-                token_scales_cls = token_scales[:, 0:1] # B, 1
-                token_scales = token_scales[:, 1:]
-            token_scales_kept = token_scales.gather(dim=1, index=index_kept) # B, N-1-num_prop
-            token_scales_prop = token_scales.gather(dim=1, index=index_prop) # B, num_prop
-            token_scales_prop = weight_prop @ token_scales_prop.unsqueeze(-1) # B, N-1-num_prop, 1
-            token_scales = token_scales_kept + alpha * token_scales_prop.squeeze(-1) # B, N-1-num_prop
-            if cls_token:
-                token_scales = torch.cat((token_scales_cls, token_scales), dim=1) # B, N-num_prop
-    else:
-        assert False, "Propagation method \'%f\' has not been supported yet." % standard
-    
-    
-    if cls_token:
-        # Step 4： concatenate the [CLS] token and generate returned value
-        x = torch.cat((x_cls, x_kept), dim=1) # B, N-num_prop, C
-    else:
-        x = x_kept
-    return x, weight, token_scales
-
-
-
-def select(weight: torch.Tensor, standard: str = "None", num_prop: int = 0, cls_token = True):
-    """
-    Select image tokens to be propagated. The [CLS] token will be ignored. 
-    ======================================================================
-    Args:
-        - weight: Tensor([B, H, N, N]): used for selecting the kept tokens. Only support the
-                                        attention map of tokens at the moment.
-        
-        - standard: str: the method applied to select the tokens
-        
-        - num_prop: int: the number of tokens to be propagated
-        
-    Return:
-        - index_kept: Tensor([B, N-1-num_prop]): the index of kept tokens 
-        
-        - index_prop: Tensor([B, num_prop]): the index of propagated tokens
-    """
-    
-    assert len(weight.shape) == 4, "Selection methods on tensors other than the attention map haven't been supported yet."
-    B, H, N1, N2 = weight.shape
-    assert N1 == N2, "Selection methods on tensors other than the attention map haven't been supported yet."
-    N = N1
-    assert num_prop >= 0, "The number of propagated/pruned tokens must be non-negative."
-    
-    if cls_token:
-        if standard == "CLSAttnMean":
-            token_rank = weight[:,:,0,1:].mean(1)
-            
-        elif standard == "CLSAttnMax":
-            token_rank = weight[:,:,0,1:].max(1)[0]
-                
-        elif standard == "IMGAttnMean":
-            token_rank = weight[:,:,:,1:].sum(-2).mean(1)
-        
-        elif standard == "IMGAttnMax":
-            token_rank = weight[:,:,:,1:].sum(-2).max(1)[0]
-                
-        elif standard == "DiagAttnMean":
-            token_rank = torch.diagonal(weight, dim1=-2, dim2=-1)[:,:,1:].mean(1)
-            
-        elif standard == "DiagAttnMax":
-            token_rank = torch.diagonal(weight, dim1=-2, dim2=-1)[:,:,1:].max(1)[0]
-            
-        elif standard == "MixedAttnMean":
-            token_rank_1 = torch.diagonal(weight, dim1=-2, dim2=-1)[:,:,1:].mean(1)
-            token_rank_2 = weight[:,:,:,1:].sum(-2).mean(1)
-            token_rank = token_rank_1 * token_rank_2
-            
-        elif standard == "MixedAttnMax":
-            token_rank_1 = torch.diagonal(weight, dim1=-2, dim2=-1)[:,:,1:].max(1)[0]
-            token_rank_2 = weight[:,:,:,1:].sum(-2).max(1)[0]
-            token_rank = token_rank_1 * token_rank_2
-            
-        elif standard == "SumAttnMax":
-            token_rank_1 = torch.diagonal(weight, dim1=-2, dim2=-1)[:,:,1:].max(1)[0]
-            token_rank_2 = weight[:,:,:,1:].sum(-2).max(1)[0]
-            token_rank = token_rank_1 + token_rank_2
-            
-        elif standard == "CosSimMean":
-            weight = weight[:,:,1:,:].mean(1)
-            weight = weight / weight.norm(dim=-1, keepdim=True)
-            token_rank = -(weight @ weight.transpose(-1, -2)).sum(-1)
-        
-        elif standard == "CosSimMax":
-            weight = weight[:,:,1:,:].max(1)[0]
-            weight = weight / weight.norm(dim=-1, keepdim=True)
-            token_rank = -(weight @ weight.transpose(-1, -2)).sum(-1)
-            
-        elif standard == "Random":
-            token_rank = torch.randn((B, N-1), device=weight.device)
-                
-        else:
-            print("Type\'", standard, "\' selection not supported.")
-            assert False
-        
-        token_rank = torch.argsort(token_rank, dim=1, descending=True) # B, N-1
-        index_kept = token_rank[:, :-num_prop]+1 # B, N-1-num_prop
-        index_prop = token_rank[:, -num_prop:]+1 # B, num_prop
-            
-    else:
-        if standard == "IMGAttnMean":
-            token_rank = weight.sum(-2).mean(1)
-        
-        elif standard == "IMGAttnMax":
-            token_rank = weight.sum(-2).max(1)[0]
-                
-        elif standard == "DiagAttnMean":
-            token_rank = torch.diagonal(weight, dim1=-2, dim2=-1).mean(1)
-            
-        elif standard == "DiagAttnMax":
-            token_rank = torch.diagonal(weight, dim1=-2, dim2=-1).max(1)[0]
-            
-        elif standard == "MixedAttnMean":
-            token_rank_1 = torch.diagonal(weight, dim1=-2, dim2=-1).mean(1)
-            token_rank_2 = weight.sum(-2).mean(1)
-            token_rank = token_rank_1 * token_rank_2
-            
-        elif standard == "MixedAttnMax":
-            token_rank_1 = torch.diagonal(weight, dim1=-2, dim2=-1).max(1)[0]
-            token_rank_2 = weight.sum(-2).max(1)[0]
-            token_rank = token_rank_1 * token_rank_2
-            
-        elif standard == "SumAttnMax":
-            token_rank_1 = torch.diagonal(weight, dim1=-2, dim2=-1).max(1)[0]
-            token_rank_2 = weight.sum(-2).max(1)[0]
-            token_rank = token_rank_1 + token_rank_2
-            
-        elif standard == "CosSimMean":
-            weight = weight.mean(1)
-            weight = weight / weight.norm(dim=-1, keepdim=True)
-            token_rank = -(weight @ weight.transpose(-1, -2)).sum(-1)
-        
-        elif standard == "CosSimMax":
-            weight = weight.max(1)[0]
-            weight = weight / weight.norm(dim=-1, keepdim=True)
-            token_rank = -(weight @ weight.transpose(-1, -2)).sum(-1)
-            
-        elif standard == "Random":
-            token_rank = torch.randn((B, N-1), device=weight.device)
-                
-        else:
-            print("Type\'", standard, "\' selection not supported.")
-            assert False
-        
-        token_rank = torch.argsort(token_rank, dim=1, descending=True) # B, N-1
-        index_kept = token_rank[:, :-num_prop] # B, N-1-num_prop
-        index_prop = token_rank[:, -num_prop:] # B, num_prop
-    return index_kept, index_prop
-
 class PoolFormerBlock(nn.Module):
     """
     Implementation of one PoolFormer block.
@@ -1212,17 +938,18 @@ class PoolFormerBlock(nn.Module):
     """
     def __init__(self, dim, pool_size=3, mlp_ratio=4., 
                  act_layer=nn.GELU, norm_layer=GroupNorm, 
-                 drop=0., drop_path=0., 
-                 use_layer_scale=True, layer_scale_init_value=1e-5):
+                 drop=0., drop_path=0., num_heads=4,
+                 use_layer_scale=True, layer_scale_init_value=1e-5, input_resolution = None, window_size = 4, shift_size = 2):
 
         super().__init__()
 
         self.norm1 = norm_layer(dim)
         #self.token_mixer = Pooling(pool_size=pool_size)
         # self.token_mixer = FNetBlock()
-        self.window_size = 4
-        self.attn_mask = None
-        self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=4)
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.input_resolution = input_resolution
+        self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, attn_drop=0.2, proj_drop=0.1)
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, 
@@ -1237,21 +964,43 @@ class PoolFormerBlock(nn.Module):
                 layer_scale_init_value * torch.ones((dim)), requires_grad=True)
             self.layer_scale_2 = nn.Parameter(
                 layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, 1, H, W))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, :, h, w] = cnt
+                    cnt += 1
 
-    def forward(self, x, weight, token_scales = None):
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x):
         B, C, H, W = x.shape
         x_windows = window_partition(x, self.window_size)
         x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
         attn_windows = self.token_mixer(x_windows, mask=self.attn_mask)
         attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
         x_attn = window_reverse(attn_windows, self.window_size, H, W)
-        index_kept, index_prop = select(x_attn, standard="MixedAttnMax", num_prop=0,
-                                            cls_token=False)
-        original_shape = x_attn.shape
-        x_attn = x_attn.view(-1, self.window_size * self.window_size, C)
-        x_attn, weight, token_scales = propagate(x_attn, weight, index_kept, index_prop, standard="GraphProp",
-                                        alpha=0.1, token_scales=token_scales, cls_token=False)
-        x_attn = x_attn.view(*original_shape)
+        if self.shift_size > 0:
+            x = torch.roll(x_attn, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = x_attn
         if self.use_layer_scale:
             x = x + self.drop_path(
                 self.layer_scale_1.unsqueeze(-1).unsqueeze(-1)