clip_vit.py

# --------------------------------------------------------
# InternVL
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
from functools import partial

import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import DropPath
from transformers import CLIPModel


def _freeze_params(module):
    for param in module.parameters():
        param.requires_grad = False


class CrossAttention(nn.Module):
    def __init__(
            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
            proj_drop=0., attn_head_dim=None, out_dim=None):
        super().__init__()
        if out_dim is None:
            out_dim = dim
        self.num_heads = num_heads
        head_dim = dim // num_heads
        if attn_head_dim is not None:
            head_dim = attn_head_dim
        all_head_dim = head_dim * self.num_heads
        self.scale = qk_scale or head_dim ** -0.5
        assert all_head_dim == dim

        self.q = nn.Linear(dim, all_head_dim, bias=False)
        self.k = nn.Linear(dim, all_head_dim, bias=False)
        self.v = nn.Linear(dim, all_head_dim, bias=False)

        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.k_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias = None
            self.k_bias = None
            self.v_bias = None

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(all_head_dim, out_dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, k=None, v=None):
        B, N, C = x.shape
        N_k = k.shape[1]
        N_v = v.shape[1]

        q_bias, k_bias, v_bias = None, None, None
        if self.q_bias is not None:
            q_bias = self.q_bias
            k_bias = self.k_bias
            v_bias = self.v_bias

        q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
        q = q.reshape(B, N, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)  # (B, N_head, N_q, dim)

        k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
        k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)

        v = F.linear(input=v, weight=self.v.weight, bias=v_bias)
        v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)

        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))  # (B, N_head, N_q, N_k)

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
        x = self.proj(x)
        x = self.proj_drop(x)

        return x


class AttentiveBlock(nn.Module):

    def __init__(self, dim, num_heads, qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, attn_head_dim=None, out_dim=None):
        super().__init__()

        self.norm1_q = norm_layer(dim)
        self.norm1_k = norm_layer(dim)
        self.norm1_v = norm_layer(dim)
        self.cross_attn = CrossAttention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
            proj_drop=drop, attn_head_dim=attn_head_dim, out_dim=out_dim)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x_q, x_kv, pos_q, pos_k, bool_masked_pos, rel_pos_bias=None):
        x_q = self.norm1_q(x_q + pos_q)
        x_k = self.norm1_k(x_kv + pos_k)
        x_v = self.norm1_v(x_kv)
        x = self.cross_attn(x_q, k=x_k, v=x_v)

        return x


class AttentionPoolingBlock(AttentiveBlock):

    def forward(self, x):
        x_q = x.mean(1, keepdim=True)
        x_kv, pos_q, pos_k = x, 0, 0
        x = super().forward(x_q, x_kv, pos_q, pos_k, bool_masked_pos=None, rel_pos_bias=None)
        x = x.squeeze(1)
        return x


class CLIPViT(nn.Module):

    def __init__(self, patch_size=14, img_size=336, pretrain_size=336, embed_dim=1024, num_heads=16,
                 mlp_ratio=4, depth=48, with_cp=True, freeze_vit=True, cls_target='cls_patch_concat',
                 num_classes=1000, pretrained=None):
        super().__init__()
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models

        self.pretrain_size = pretrain_size
        self.img_size = img_size
        self.patch_size = patch_size
        self.cls_target = cls_target
        self.depth = depth
        self.mlp_ratio = mlp_ratio
        self.with_cp = with_cp

        model = CLIPModel.from_pretrained(pretrained)
        model.post_layernorm = nn.Identity()
        self.model = model.vision_model

        if freeze_vit:
            _freeze_params(self)

        if cls_target == 'cls_patch_concat':
            self.norm = nn.SyncBatchNorm(embed_dim * 2, eps=1e-6)
            self.head = nn.Linear(embed_dim * 2, num_classes) if num_classes > 0 else nn.Identity()
        elif cls_target == 'attention_pooling':
            self.attn_pooling = AttentionPoolingBlock(
                dim=embed_dim, num_heads=num_heads, qkv_bias=True, qk_scale=None,
                drop=0., attn_drop=0.0, norm_layer=partial(nn.LayerNorm, eps=1e-5), out_dim=embed_dim)
            self.norm = nn.SyncBatchNorm(embed_dim, eps=1e-6)
            self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
        else:
            raise NotImplementedError

        if type(self.head) != nn.Identity:
            self.head.weight.data.normal_(mean=0.0, std=0.01)
            self.head.bias.data.zero_()

    @property
    def dtype(self):
        return self.model.embeddings.patch_embedding.weight.dtype

    def forward_features(self, x):
        x = x.type(self.dtype)
        x = self.model(x)
        x = x.last_hidden_state
        return x

    def forward(self, x):
        x = self.forward_features(x)
        if self.cls_target == 'cls_patch_concat':
            x = torch.cat((x[:, 0, :], x[:, 1:, :].mean(dim=1)), dim=-1)
        elif self.cls_target == 'attention_pooling':
            x = self.attn_pooling(x)
        else:
            raise NotImplementedError
        x = self.norm(x)
        x = self.head(x)
        return x

    @torch.jit.ignore
    def lr_decay_keywords(self, decay_ratio=0.95):
        lr_ratios = {}

        # layers
        for idx in range(self.depth):
            tag = 'layers.{}.'.format(idx)
            decay = 1.0 * (decay_ratio ** (self.depth - idx))
            lr_ratios[tag] = decay

        # patch_embedding
        lr_ratios['patch_embedding'] = 1.0 * (decay_ratio ** (self.depth + 1))
        lr_ratios['position_embedding'] = 1.0 * (decay_ratio ** (self.depth + 1))
        lr_ratios['pre_layrnorm'] = 1.0 * (decay_ratio ** (self.depth + 1))

        return lr_ratios