Implementing Swin Transformer Blocks from Scratch in PyTorch

Standard Vision Transformers compute global self-attention across every patch, which scales quadratically with image resolution. Swin Transformers fix this by computing attention inside local windows and shifting those windows between layers, giving the model a way to exchange information across window boundaries without paying the full quadratic cost.

Windowed Self-Attention

Instead of attending over all N patches, Swin partitions the feature map into non-overlapping M x M windows and computes standard multi-head self-attention independently within each window. This drops the attention complexity from O(N^2) to O(N * M^2) — linear in the number of patches for a fixed window size.

class WindowAttention(nn.Module):
    def __init__(self, dim, window_size, num_heads):
        super().__init__()
        self.window_size = window_size
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=True)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x):
        B_, N, C = x.shape
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        return self.proj(x)

The Shifted Window

Local windows alone can't model relationships between patches in different windows. Swin alternates between regular and shifted window partitioning every other block — shifting by (M//2, M//2) pixels — so that a patch near a window boundary in one layer ends up in the interior of a window in the next layer, and attention naturally propagates information across the original boundary.

The practical trick is doing this efficiently with a single torch.roll plus a masked attention pass, rather than physically re-padding windows:

shifted_x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
# attention computed on shifted_x with an attn_mask that
# zeroes out cross-region terms introduced by the roll

Why It Matters in Production

For dense prediction tasks — segmentation, detection backbones — the linear complexity is what makes Swin viable as a drop-in ResNet replacement at real input resolutions, without the memory blowup a plain ViT would hit. This is the backbone family behind several of the edge-inference classification stacks in the flagship systems on this site.

This post is a technical walkthrough for the site's engineering blog and will be expanded with benchmark numbers from the training runs it references.