ta012/SSLAM_AS2M_Finetuned

1import torch2import torch.nn as nn3import torch.nn.functional as F4import numpy as np5from timm.models.layers import to_2tuple6 7class PatchEmbed_new(nn.Module):8    """ Flexible Image to Patch Embedding9    """10    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, stride=16):11        super().__init__()12        img_size = to_2tuple(img_size)13        patch_size = to_2tuple(patch_size)14        stride = to_2tuple(stride)15        16        self.img_size = img_size17        self.patch_size = patch_size18 19        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride) # with overlapped patches20 21    def forward(self, x):22        x = self.proj(x)23        x = x.flatten(2).transpose(1, 2)24        return x25 26 27def get_2d_sincos_pos_embed_flexible(embed_dim, grid_size, cls_token=False):28    """29    grid_size: int of the grid height and width30    return:31    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)32    """33    grid_h = np.arange(grid_size[0], dtype=np.float32)34    grid_w = np.arange(grid_size[1], dtype=np.float32)35    grid = np.meshgrid(grid_w, grid_h)  # here w goes first36    grid = np.stack(grid, axis=0)37 38    grid = grid.reshape([2, 1, grid_size[0], grid_size[1]])39    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)40    if cls_token:41        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)42    return pos_embed43 44 45def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):46    assert embed_dim % 2 == 047 48    # use half of dimensions to encode grid_h49    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)50    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)51 52    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)53    return emb54 55 56def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):57    """58    embed_dim: output dimension for each position59    pos: a list of positions to be encoded: size (M,)60    out: (M, D)61    """62    assert embed_dim % 2 == 063    omega = np.arange(embed_dim // 2, dtype=np.float32)64    omega /= embed_dim / 2.065    omega = 1.0 / 10000 ** omega  # (D/2,)66 67    pos = pos.reshape(-1)  # (M,)68    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product69 70    emb_sin = np.sin(out)  # (M, D/2)71    emb_cos = np.cos(out)  # (M, D/2)72 73    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)74    return emb75 76 77class FixedPositionalEncoder(nn.Module):78    def __init__(self, pos_embed):79        super().__init__()80        self.positions = pos_embed81 82    def forward(self, x, padding_mask): 83        return self.positions84 85 86class AltBlock(nn.Module):87    def __init__(88        self,89        dim,90        num_heads,91        mlp_ratio=4.0,92        qkv_bias=False,93        qk_scale=None,94        drop=0.0,95        attn_drop=0.0,96        mlp_drop=0.0,97        post_mlp_drop=0.0,98        drop_path=0.0,99        act_layer=nn.GELU,100        norm_layer=nn.LayerNorm,101        layer_norm_first=True,102        ffn_targets=False,103        cosine_attention=False,104    ):105        super().__init__()106 107        self.layer_norm_first = layer_norm_first108        self.ffn_targets = ffn_targets109 110        from timm.models.vision_transformer import DropPath, Mlp111 112        self.norm1 = norm_layer(dim)113        self.attn = AltAttention(114            dim,115            num_heads=num_heads,116            qkv_bias=qkv_bias,117            qk_scale=qk_scale,118            attn_drop=attn_drop,119            proj_drop=drop,120            cosine_attention=cosine_attention,121        )122 123        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()124        self.norm2 = norm_layer(dim)125        mlp_hidden_dim = int(dim * mlp_ratio)126        self.mlp = Mlp(127            in_features=dim,128            hidden_features=mlp_hidden_dim,129            act_layer=act_layer,130            drop=mlp_drop,131        )132        self.post_mlp_dropout = nn.Dropout(post_mlp_drop, inplace=False)133 134    def forward(self, x, padding_mask=None, alibi_bias=None):135        if self.layer_norm_first:136            x = x + self.drop_path(self.attn(self.norm1(x), padding_mask, alibi_bias))137            r = x = self.mlp(self.norm2(x))138            t = x139            x = r + self.drop_path(self.post_mlp_dropout(x))140            if not self.ffn_targets:141                t = x142        else:143            x = x + self.drop_path(self.attn(x, padding_mask, alibi_bias))144            r = x = self.norm1(x)145            x = self.mlp(x)146            t = x147            x = self.norm2(r + self.drop_path(self.post_mlp_dropout(x)))148            if not self.ffn_targets:149                t = x150 151        return x, t152 153 154class AltAttention(nn.Module):155    def __init__(156        self,157        dim,158        num_heads=8,159        qkv_bias=False,160        qk_scale=None,161        attn_drop=0.0,162        proj_drop=0.0,163        cosine_attention=False,164    ):165        super().__init__()166        self.num_heads = num_heads167        head_dim = dim // num_heads168        self.scale = qk_scale or head_dim ** -0.5169 170        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)171        self.attn_drop = nn.Dropout(attn_drop)172        self.proj = nn.Linear(dim, dim)173        self.proj_drop = nn.Dropout(proj_drop)174 175        self.cosine_attention = cosine_attention176 177        if cosine_attention:178            self.logit_scale = nn.Parameter(179                torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True180            )181 182    def forward(self, x, padding_mask=None, alibi_bias=None):183        B, N, C = x.shape184        qkv = (185            self.qkv(x)186            .reshape(B, N, 3, self.num_heads, C // self.num_heads)187            .permute(2, 0, 3, 1, 4)  # qkv x B x H x L x D188        )189        q, k, v = (190            qkv[0],191            qkv[1],192            qkv[2],193        )  # make torchscript happy (cannot use tensor as tuple)194 195        dtype = q.dtype196 197        if self.cosine_attention:198            # cosine attention199            attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)200            logit_scale = torch.clamp(201                self.logit_scale, max=torch.log(torch.tensor(1.0 / 0.01))202            ).exp()203            attn = attn * logit_scale204        else:205            q = q * self.scale206            attn = q @ k.transpose(-2, -1)207 208        if alibi_bias is not None:209            attn = attn.type_as(alibi_bias)210            attn[:, : alibi_bias.size(1)] += alibi_bias211 212        if padding_mask is not None and padding_mask.any():213            attn = attn.masked_fill(214                padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),215                float("-inf"),216            )217 218        attn = attn.softmax(dim=-1, dtype=torch.float32).to(dtype=dtype)219        attn = self.attn_drop(attn)220        x = (attn @ v).transpose(1, 2)  #221        x = x.reshape(B, N, C)222        x = self.proj(x)223        x = self.proj_drop(x)224        return x