Upload vitGPT.py

vitGPT.py

@@ -0,0 +1,349 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from timm import create_model, list_models
+from types import SimpleNamespace
+from transformers import GPT2LMHeadModel, GPT2TokenizerFast, get_linear_schedule_with_warmup
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+from pathlib import Path
+from sklearn.model_selection import train_test_split
+from torch.cuda.amp import GradScaler, autocast
+from tqdm.auto import tqdm
+import gc
+import json
+
+class GPT2Attention(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+        self.embed_dim = config.embed_dim
+        self.n_heads = config.num_heads
+        assert self.embed_dim % self.n_heads == 0, 'embedding dimension by be divisible by number of heads'
+        self.head_size = self.embed_dim // self.n_heads
+        self.seq_len = config.seq_len
+
+        self.c_attn = nn.Linear(self.embed_dim, self.head_size * self.n_heads * 3,bias=True)
+        self.scale = self.head_size ** -0.5
+
+        self.register_buffer('mask',torch.tril(torch.ones(1,1,self.seq_len,self.seq_len)))
+
+        self.c_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+        self.attn_dropout = nn.Dropout(config.attention_dropout)
+        self.resid_dropout = nn.Dropout(config.residual_dropout)
+
+
+    def forward(self, x):
+        b,t,c = x.shape
+        # q,k,v shape individually: batch_size x seq_len x embed_dim
+        # we know that qk_t = q x k_t, where q=bxtxhead_dim, k_t=bxhead_timxt
+        q,k,v = self.c_attn(x).chunk(3,dim=-1)
+        q = q.view(b,t,self.n_heads,self.head_size).permute(0,2,1,3) # batch x n_heads x seq_len x head_dim
+        k = k.view(b,t,self.n_heads,self.head_size).permute(0,2,1,3)
+        v = v.view(b,t,self.n_heads,self.head_size).permute(0,2,1,3)
+
+        qk_t = (q@k.transpose(-2,-1)) * self.scale
+        qk_t = qk_t.masked_fill(self.mask[:,:,:t,:t]==0,float('-inf'))
+        qk_t = F.softmax(qk_t,dim=-1)
+        weights = self.attn_dropout(qk_t)
+
+        attention = weights @ v # batch x n_heads x t x head_size
+        attention = attention.permute(0,2,1,3).contiguous().view(b,t,c) # batch x t x embed_dim
+
+        out = self.c_proj(attention)
+        out = self.resid_dropout(out)
+
+        return out
+
+class GPT2CrossAttention(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+        self.embed_dim = config.embed_dim
+        self.n_heads = config.num_heads
+        assert self.embed_dim % self.n_heads == 0, 'embedding dimension by be divisible by number of heads'
+        self.head_size = self.embed_dim // self.n_heads
+        self.seq_len = config.seq_len
+
+        self.q = nn.Linear(self.embed_dim,self.embed_dim)
+        self.k = nn.Linear(self.embed_dim,self.embed_dim)
+        self.v = nn.Linear(self.embed_dim,self.embed_dim)
+        self.scale = self.head_size ** -0.5
+
+        self.c_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+        self.attn_dropout = nn.Dropout(config.attention_dropout)
+        self.resid_dropout = nn.Dropout(config.residual_dropout)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+
+
+    def forward(self, q,k,v):
+        b,t,c = q.shape
+
+        q = self.q(q)
+        k = self.k(k)
+        v = self.v(v)
+
+        q = q.view(b,q.size(1),self.n_heads,self.head_size).permute(0,2,1,3) # batch x n_heads x seq_len x head_dim
+        k = k.view(b,k.size(1),self.n_heads,self.head_size).permute(0,2,1,3)
+        v = v.view(b,v.size(1),self.n_heads,self.head_size).permute(0,2,1,3)
+
+        qk_t = (q@k.transpose(-2,-1)) * self.scale
+        qk_t = F.softmax(qk_t,dim=-1)
+        weights = self.attn_dropout(qk_t)
+
+        attention = weights @ v # batch x n_heads x t x head_size
+        attention = attention.permute(0,2,1,3).contiguous().view(b,t,c) # batch x t x embed_dim
+
+        out = self.c_proj(attention)
+        out = self.resid_dropout(out)
+
+        return out
+
+
+class GPT2MLP(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+        self.embed_dim = config.embed_dim
+        self.mlp_ratio = config.mlp_ratio
+        self.mlp_dropout = config.mlp_dropout
+
+        self.c_fc = nn.Linear(self.embed_dim,self.embed_dim*self.mlp_ratio)
+        self.c_proj = nn.Linear(self.embed_dim*self.mlp_ratio,self.embed_dim)
+        self.act = nn.GELU()
+        self.dropout = nn.Dropout(self.mlp_dropout)
+
+    def forward(self,x):
+        x = self.c_fc(x)
+        x = self.act(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class GPT2Block(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+        self.embed_dim = config.embed_dim
+        self.ln_1 = nn.LayerNorm(self.embed_dim)
+        self.attn = GPT2Attention(config)
+        self.ln_2 = nn.LayerNorm(self.embed_dim)
+        self.mlp = GPT2MLP(config)
+        self.ln_3 = nn.LayerNorm(self.embed_dim)
+        self.cross_attn = GPT2CrossAttention(config)
+
+    def forward(self,x,enc_out):
+        x = x+self.attn(self.ln_1(x))
+        x = x+self.cross_attn(self.ln_2(x),enc_out,enc_out)
+        x = x+self.mlp(self.ln_3(x))
+        return x
+
+
+
+class VisionGPT2Model(nn.Module):
+    def __init__(self,config):
+        super().__init__()
+
+        self.config = config
+        print(torch.cuda.is_available())
+        vit = create_model('vit_base_patch16_224',pretrained=True,num_classes=0)
+        self.patch_embed = vit.patch_embed
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = vit.cls_token
+        embed_len = num_patches + vit.num_prefix_tokens
+        self.pos_embed = vit.pos_embed
+        self.pos_drop = nn.Dropout(p=0.)
+
+        self.blocks = nn.ModuleList([vit.blocks[i] for i in range(config.depth)])
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size,config.embed_dim),
+            wpe = nn.Embedding(config.seq_len,config.embed_dim),
+            drop = nn.Dropout(config.emb_dropout),
+            h = nn.ModuleList([GPT2Block(config) for _ in range(config.depth)]),
+            ln_f = nn.LayerNorm(config.embed_dim)
+        ))
+        self.lm_head = nn.Linear(config.embed_dim,config.vocab_size,bias=False)
+        self.transformer.wte.weight = self.lm_head.weight
+
+    def _pos_embed(self,x):
+        pos_embed = self.pos_embed
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + pos_embed
+        return self.pos_drop(x)
+
+    def pretrained_layers_trainable(self,trainable=False):
+        layers = [
+            self.cls_token, self.patch_embed, self.pos_embed, self.blocks,
+            self.transformer.wte, self.transformer.wpe,
+            self.transformer.ln_f, self.lm_head
+        ]
+        gpt_layers = [[
+            self.transformer.h[i].ln_1,self.transformer.h[i].ln_2,
+            self.transformer.h[i].attn,self.transformer.h[i].mlp
+        ] for i in range(self.config.depth)]
+        for l in gpt_layers:
+            layers.extend(l)
+
+        for layer in layers:
+            if not isinstance(layer,nn.Parameter):
+                for p in layer.parameters():
+                    p.requires_grad = trainable
+            else:
+                layer.requires_grad = trainable
+
+        total_frozen_params = sum([p.numel() for p in self.parameters() if not p.requires_grad])
+        print(f'{total_frozen_params=}')
+
+    def unfreeze_gpt_layers(self,):
+        gpt_layers = [[
+            self.transformer.h[i].ln_1,self.transformer.h[i].ln_2,
+            self.transformer.h[i].attn,self.transformer.h[i].mlp
+        ] for i in range(self.config.depth)]
+        flatten = []
+        for l in gpt_layers:
+            flatten.extend(l)
+
+        for layer in flatten:
+            if not isinstance(layer,nn.Parameter):
+                for p in layer.parameters():
+                    p.requires_grad = True
+            else:
+                layer.requires_grad = True
+
+    @classmethod
+    def from_pretrained(self,config):
+        model = VisionGPT2Model(config)
+        sd = model.state_dict()
+        keys = sd.keys()
+        ignore_matches = ['blocks.','cross_attn.','ln_3','cls_token','pos_embed','patch_embed.','.attn.mask']
+        vit_keys = [key for key in keys if any(match in key for match in ignore_matches)]
+        gpt_keys = [key for key in keys if key not in vit_keys]
+
+        gpt2_small = GPT2LMHeadModel.from_pretrained('gpt2')
+        sd_hf = gpt2_small.state_dict()
+        hf_keys = sd_hf.keys()
+        hf_keys = [k for k in hf_keys if not k.endswith('.attn.masked_bias')]
+        hf_keys = [k for k in hf_keys if not k.endswith('.attn.bias')]
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+
+        for k in hf_keys:
+            if any(match in k for match in ignore_matches):
+                continue
+            if any(k.endswith(w) for w in transposed):
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        model.load_state_dict(sd)
+
+        return model
+
+    def forward(self,image,input_ids,labels=None):
+
+        image = self.patch_embed(image)
+        image = self._pos_embed(image)
+
+        token_embeddings = self.transformer.wte(input_ids) # batch x seq_len
+        pos_embs = torch.arange(0, input_ids.size(1), device=self.config.device)
+        positional_embeddings = self.transformer.wpe(pos_embs)
+        input_ids = self.transformer.drop(token_embeddings+positional_embeddings)
+
+        for i in range(self.config.depth):
+            image = self.blocks[i](image)
+            input_ids = self.transformer.h[i](input_ids, image)
+
+        input_ids = self.transformer.ln_f(input_ids)
+
+        if labels is not None:
+            lm_logits = self.lm_head(input_ids)
+            loss = F.cross_entropy(lm_logits.view(-1, lm_logits.shape[-1]), labels.view(-1))
+            return loss
+
+        lm_logits = self.lm_head(input_ids[:,[-1],:])
+        return lm_logits
+
+    def generate(self,image,sequence,tokenizer,max_tokens=50,temperature=1.0,deterministic=False):
+        for _ in range(max_tokens):
+            out = self(image,sequence)
+            out = out[:,-1,:] / temperature
+            probs = F.softmax(out,dim=-1)
+            if deterministic:
+                next_token = torch.argmax(probs,dim=-1,keepdim=True)
+            else:
+                next_token = torch.multinomial(probs,num_samples=1)
+            sequence = torch.cat([sequence,next_token],dim=1)
+            if next_token.item() == tokenizer.eos_token_id:
+                break
+
+        return sequence.cpu().flatten()
+
+
+model_config = SimpleNamespace(
+    vocab_size = 50_257,
+    embed_dim = 768,
+    num_heads = 12,
+    seq_len = 1024,
+    depth = 12,
+    attention_dropout = 0.1,
+    residual_dropout = 0.1,
+    mlp_ratio = 4,
+    mlp_dropout = 0.1,
+    emb_dropout = 0.1,
+    device='cpu'
+)
+
+
+
+model = VisionGPT2Model.from_pretrained(model_config)  
+model.load_state_dict(torch.load("captioner.pt", map_location='cpu')) # Use 'cuda' if you have a GPU
+model.eval()  # Set the model to evaluation mode
+
+
+def generate_caption(image,max_tokens=50,temperature=0.9,deterministic=True):
+    tokenizer = GPT2TokenizerFast.from_pretrained('gpt2')
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+    gen_tfms = A.Compose([
+            A.Resize(224,224),
+            A.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5],always_apply=True),
+            ToTensorV2()
+        ])
+
+    image = Image.open(image)
+    image = np.array(image)
+    image = gen_tfms(image=image)['image']
+    image = image.unsqueeze(0)
+    sequence = torch.ones(1,1).long() * tokenizer.bos_token_id
+
+    caption = model.generate(
+        image,
+        sequence,
+        tokenizer,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        deterministic=deterministic, 
+        
+    )
+    caption = tokenizer.decode(caption.numpy(),skip_special_tokens=True)
+    print(caption)
+    return caption
+
+image = "/Users/jkottu/Desktop/image-captioning-chest-xrays/sample_images/CXR191_IM-0591-1001.png"
+generate_caption(image)