swin_transformer_rope.py

"""
This code was originally obtained from:
https://github.com/microsoft/Swin-Transformer
and
https://github.com/meta-llama/codellama/blob/main/llama/model.py
"""
import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint
import torch.nn.functional as F
from typing import Any, Optional, Tuple
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from functools import partial

from SwinTransformer.models.swin_transformer import SwinTransformer, Mlp, SwinTransformerBlock, WindowAttention, BasicLayer
from SwinTransformer.models.swin_transformer import window_partition, window_reverse, PatchMerging
from SwinTransformer.models.swin_transformer import WindowProcess, WindowProcessReverse

## RoPE functions

def init_t_xy(end_x: int, end_y: int, zero_center=False):
    t = torch.arange(end_x * end_y, dtype=torch.float32)
    t_x = (t % end_x).float()
    t_y = torch.div(t, end_x, rounding_mode='floor').float()
    
    return t_x, t_y

def init_random_2d_freqs(head_dim: int, num_heads: int, theta: float = 10.0, rotate: bool = True):
    freqs_x = []
    freqs_y = []
    theta = theta
    mag = 1 / (theta ** (torch.arange(0, head_dim, 4)[: (head_dim // 4)].float() / head_dim))
    for i in range(num_heads):
        angles = torch.rand(1) * 2 * torch.pi if rotate else torch.zeros(1)
        fx = torch.cat([mag * torch.cos(angles), mag * torch.cos(torch.pi/2 + angles)], dim=-1)
        fy = torch.cat([mag * torch.sin(angles), mag * torch.sin(torch.pi/2 + angles)], dim=-1)
        freqs_x.append(fx)
        freqs_y.append(fy)
    freqs_x = torch.stack(freqs_x, dim=0)
    freqs_y = torch.stack(freqs_y, dim=0)
    freqs = torch.stack([freqs_x, freqs_y], dim=0)
    return freqs

def compute_cis(freqs: torch.Tensor, t_x: torch.Tensor, t_y: torch.Tensor):
    N = t_x.shape[0]
    # No float 16 for this range
    with torch.cuda.amp.autocast(enabled=False):
        freqs_x = (t_x.unsqueeze(-1) @ freqs[0].unsqueeze(-2))
        freqs_y = (t_y.unsqueeze(-1) @ freqs[1].unsqueeze(-2))
        freqs_cis = torch.polar(torch.ones_like(freqs_x), freqs_x + freqs_y)
        
    return freqs_cis


def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    assert 0 <= 1 < ndim
    # assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
    if freqs_cis.shape == (x.shape[-2], x.shape[-1]):
        shape = [d if i >= ndim-2 else 1 for i, d in enumerate(x.shape)]
    elif freqs_cis.shape == (x.shape[-3], x.shape[-2], x.shape[-1]):
        shape = [d if i >= ndim-3 else 1 for i, d in enumerate(x.shape)]
        
    return freqs_cis.view(*shape)

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)

class RoPEWindowAttention(WindowAttention):

    def __init__(self, *args, rope_theta=10.0, rope_mixed=True, use_rpb=False, **kwargs):

        super().__init__(*args, **kwargs)

        self.rope_mixed = rope_mixed
        self.use_rpb = use_rpb
        if not self.use_rpb:
            self.relative_position_bias_table = None
            self.relative_position_index = None
        
        t_x, t_y = init_t_xy(end_x=self.window_size[1], end_y=self.window_size[0])
        self.register_buffer('rope_t_x', t_x)
        self.register_buffer('rope_t_y', t_y)

        freqs = init_random_2d_freqs(
            head_dim=self.dim // self.num_heads, num_heads=self.num_heads, theta=rope_theta, 
            rotate=self.rope_mixed
        )
        if self.rope_mixed:
            self.rope_freqs = nn.Parameter(freqs, requires_grad=True)
        else:
            self.register_buffer('rope_freqs', freqs)
            freqs_cis = compute_cis(self.rope_freqs, self.rope_t_x, self.rope_t_y)
            self.rope_freqs_cis = freqs_cis
        
    def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        q = q * self.scale

        if self.rope_mixed:
            freqs_cis = compute_cis(self.rope_freqs, self.rope_t_x, self.rope_t_y)
        else:
            freqs_cis = self.rope_freqs_cis.to(x.device)
        q, k = apply_rotary_emb(q, k, freqs_cis)
        
        attn = (q @ k.transpose(-2, -1))

        if self.use_rpb:
            relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
            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)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class RoPESwinTransformerBlock(SwinTransformerBlock):

    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 fused_window_process=False,
                 rope_theta=10.0, rope_mixed=True, use_rpb=False):
        super().__init__(
            dim, input_resolution, num_heads, window_size=window_size, shift_size=shift_size,
            mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop,
            attn_drop=attn_drop, drop_path=drop_path, act_layer=act_layer, norm_layer=norm_layer,
            fused_window_process=fused_window_process
        )

        self.attn = RoPEWindowAttention(
            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
            rope_theta=rope_theta, rope_mixed=rope_mixed, use_rpb=use_rpb
        )


class RoPEBasicLayer(BasicLayer):

    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
                 fused_window_process=False,
                 rope_theta=10.0, rope_mixed=True, use_rpb=False):

        super().__init__(
            dim=dim, input_resolution=input_resolution, depth=depth, num_heads=num_heads,
            window_size=window_size, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
            drop=drop, attn_drop=attn_drop, drop_path=drop_path, norm_layer=norm_layer,
            downsample=downsample, use_checkpoint=use_checkpoint, fused_window_process=fused_window_process
        )

        # build blocks
        self.blocks = nn.ModuleList([
            RoPESwinTransformerBlock(
                dim=dim, input_resolution=input_resolution,
                num_heads=num_heads, window_size=window_size,
                shift_size=0 if (i % 2 == 0) else window_size // 2,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop, attn_drop=attn_drop,
                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
                norm_layer=norm_layer,
                fused_window_process=fused_window_process,
                rope_theta=rope_theta, rope_mixed=rope_mixed, use_rpb=use_rpb
            )
            for i in range(depth)])


class RoPESwinTransformer(SwinTransformer):
    
    def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
                 use_checkpoint=False, fused_window_process=False,
                 rope_theta=10.0, rope_mixed=True, use_rpb=False,
                 **kwargs):

        super().__init__(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, num_classes=num_classes,
            embed_dim=embed_dim, depths=depths, num_heads=num_heads,
            window_size=window_size, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
            qk_scale=qk_scale, drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
            drop_path_rate=drop_path_rate, norm_layer=norm_layer, ape=ape,
            patch_norm=patch_norm, use_checkpoint=use_checkpoint, **kwargs
        )

        # absolute position embedding
        self.ape = False
        self.absolute_pos_embed = None

        patches_resolution = self.patch_embed.patches_resolution

        # stochastic depth
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
        
        # build layers
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = RoPEBasicLayer(
                dim=int(embed_dim * 2 ** i_layer),
                input_resolution=(patches_resolution[0] // (2 ** i_layer),
                                  patches_resolution[1] // (2 ** i_layer)),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=self.mlp_ratio,
                qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate,
                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                norm_layer=norm_layer,
                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
                use_checkpoint=use_checkpoint,
                fused_window_process=fused_window_process,
                rope_theta=rope_theta, rope_mixed=rope_mixed, use_rpb=use_rpb
            )
            self.layers.append(layer)

        self.apply(self._init_weights)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'absolute_pos_embed'}

    @torch.jit.ignore
    def no_weight_decay_keywords(self):
        return {'rope_freqs', 'relative_position_bias_table'}

def hf_checkpoint_load(model_name):
    try:
        from huggingface_hub import hf_hub_download

        ckpt_path = hf_hub_download(
            repo_id="naver-ai/" + model_name, filename= "pytorch_model.bin"
        )
        checkpoint = torch.load(ckpt_path, map_location='cpu')
    except:
        _HF_URL = "https://huggingface.co/naver-ai/" + model_name + "/resolve/main/pytorch_model.bin"
        checkpoint = torch.hub.load_state_dict_from_url(_HF_URL, map_location='cpu')

    state_dict = checkpoint['model']
    
    # delete rope_t since we always re-init it
    rope_t_keys = [k for k in state_dict.keys() if "rope_t_" in k]
    for k in rope_t_keys:
        del state_dict[k]

    # delete relative_position_index since we always re-init it
    relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k]
    for k in relative_position_index_keys:
        del state_dict[k]

    # delete relative_coords_table since we always re-init it
    relative_position_index_keys = [k for k in state_dict.keys() if "relative_coords_table" in k]
    for k in relative_position_index_keys:
        del state_dict[k]

    # delete attn_mask since we always re-init it
    attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k]
    for k in attn_mask_keys:
        del state_dict[k]
            
    return state_dict

def swin_rope_mixed_tiny_patch4_window7_224(pretrained=False, img_size=224):
    window_size = img_size // 32
    model = RoPESwinTransformer(
        img_size=img_size, patch_size=4, embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
        window_size=window_size, mlp_ratio=4, rope_theta=10.0, rope_mixed=True, use_rpb=False
    )
    if pretrained:
        state_dict = hf_checkpoint_load("swin_rope_mixed_tiny_patch4_window7_224")
        model.load_state_dict(state_dict, strict=False)
    return model

def swin_rope_mixed_small_patch4_window7_224(pretrained=False, img_size=224):
    window_size = img_size // 32
    model = RoPESwinTransformer(
        img_size=img_size, patch_size=4, embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24],
        window_size=window_size, mlp_ratio=4, rope_theta=10.0, rope_mixed=True, use_rpb=False
    )
    if pretrained:
        state_dict = hf_checkpoint_load("swin_rope_mixed_small_patch4_window7_224")
        model.load_state_dict(state_dict, strict=False)
    return model

def swin_rope_mixed_base_patch4_window7_224(pretrained=False, img_size=224):
    window_size = img_size // 32
    model = RoPESwinTransformer(
        img_size=img_size, patch_size=4, embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32],
        window_size=window_size, mlp_ratio=4, rope_theta=10.0, rope_mixed=True, use_rpb=False
    )
    if pretrained:
        state_dict = hf_checkpoint_load("swin_rope_mixed_base_patch4_window7_224")
        model.load_state_dict(state_dict, strict=False)
    return model

def swin_rope_axial_tiny_patch4_window7_224(pretrained=False, img_size=224):
    window_size = img_size // 32
    model = RoPESwinTransformer(
        img_size=img_size, patch_size=4, embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
        window_size=window_size, mlp_ratio=4, rope_theta=50.0, rope_mixed=False, use_rpb=False
    )
    if pretrained:
        state_dict = hf_checkpoint_load("swin_rope_axial_tiny_patch4_window7_224")
        model.load_state_dict(state_dict, strict=False)
    return model

def swin_rope_axial_small_patch4_window7_224(pretrained=False, img_size=224):
    window_size = img_size // 32
    model = RoPESwinTransformer(
        img_size=img_size, patch_size=4, embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24],
        window_size=window_size, mlp_ratio=4, rope_theta=50.0, rope_mixed=False, use_rpb=False
    )
    if pretrained:
        state_dict = hf_checkpoint_load("swin_rope_axial_small_patch4_window7_224")
        model.load_state_dict(state_dict, strict=False)
    return model

def swin_rope_axial_base_patch4_window7_224(pretrained=False, img_size=224):
    window_size = img_size // 32
    model = RoPESwinTransformer(
        img_size=img_size, patch_size=4, embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32],
        window_size=window_size, mlp_ratio=4, rope_theta=50.0, rope_mixed=False, use_rpb=False
    )
    if pretrained:
        state_dict = hf_checkpoint_load("swin_rope_axial_base_patch4_window7_224")
        model.load_state_dict(state_dict, strict=False)
    return model