fix dynamic padding edge case

ain-soph · ain-soph · commit 1c7579cb1bd7 · 2022-07-12T18:43:30.000-07:00
diff --git a/torchvision/models/swin_transformer.py b/torchvision/models/swin_transformer.py
@@ -1,6 +1,6 @@
 import math
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -143,18 +143,21 @@ def shifted_window_attention(
         Tensor[N, H, W, C]: The output tensor after shifted window attention.
     """
     B, H, W, C = input.shape
+
+    # If window size is larger than feature size, there is no need to shift window
+    if window_size[0] >= H:
+        shift_size[0] = 0
+        window_size[0] = H
+    if window_size[1] >= W:
+        shift_size[1] = 0
+        window_size[1] = W
+
     # pad feature maps to multiples of window size
     pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
     pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
     x = F.pad(input, (0, 0, 0, pad_r, 0, pad_b))
     _, pad_H, pad_W, _ = x.shape
 
-    # If window size is larger than feature size, there is no need to shift window
-    if window_size[0] >= pad_H:
-        shift_size[0] = 0
-    if window_size[1] >= pad_W:
-        shift_size[1] = 0
-
     # cyclic shift
     if sum(shift_size) > 0:
         x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2))
@@ -479,7 +482,7 @@ class SwinTransformer(nn.Module):
         embed_dim (int): Patch embedding dimension.
         depths (List(int)): Depth of each Swin Transformer layer.
         num_heads (List(int)): Number of attention heads in different layers.
-        window_size (List[int]): Window size.
+        window_size (int, List[int]): Window size.
         mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
         stochastic_depth_prob (float): Stochastic depth rate. Default: 0.1.
         num_classes (int): Number of classes for classification head. Default: 1000.
