| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from transformers import GPT2LMHeadModel, GPT2TokenizerFast |
| | from timm import create_model |
| | from types import SimpleNamespace |
| |
|
| |
|
| | tokenizer = GPT2TokenizerFast.from_pretrained('gpt2') |
| | tokenizer.pad_token = tokenizer.eos_token |
| |
|
| | class GPT2Attention(nn.Module): |
| | def __init__(self,config:SimpleNamespace): |
| | super(GPT2Attention,self).__init__() |
| | self.embed_dim = config.embed_dim |
| | self.n_heads = config.num_heads |
| | assert self.embed_dim % self.n_heads == 0, "embedding dim must be divisible by num 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.embed_dim*3) |
| | 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) |
| | self.attn_dropout = nn.Dropout(config.attention_dropout) |
| | self.resid_dropout = nn.Dropout(config.residual_dropout) |
| | |
| | def forward(self,x:torch.Tensor)-> torch.Tensor: |
| | b,t,c = x.shape |
| | |
| | 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) |
| | 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 |
| | attention = attention.permute(0,2,1,3).contiguous().view(b,t,c) |
| | |
| | out = self.c_proj(attention) |
| | return self.resid_dropout(out) |
| |
|
| | class GPT2CrossAttention(nn.Module): |
| | def __init__(self,config:SimpleNamespace): |
| | super(GPT2CrossAttention,self).__init__() |
| | self.embed_dim = config.embed_dim |
| | self.n_heads = config.num_heads |
| | assert self.embed_dim %self.n_heads == 0, "embedding dim must be divisible by num 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) |
| | 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): |
| | nn.init.normal_(module.weight,mean=0.0,std=0.02) |
| | |
| | if module.bias is not None: |
| | nn.init.zeros_(module.bias) |
| | |
| | def forward(self,q:torch.Tensor,k:torch.Tensor,v:torch.Tensor)->torch.Tensor: |
| | b,t,c = q.shape |
| | |
| | q,k,v = self.q(q),self.k(k),self.v(v) |
| | |
| | q = q.view(b,q.size(1),self.n_heads,self.head_size).permute(0,2,1,3) |
| | 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 |
| | attention = attention.permute(0,2,1,3).contiguous().view(b,t,c) |
| | |
| | out = self.c_proj(attention) |
| | return self.resid_dropout(out) |
| | |
| | class GPT2MLP(nn.Module): |
| | def __init__(self,config:SimpleNamespace): |
| | 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:torch.Tensor)->torch.Tensor: |
| | x = self.c_fc(x) |
| | x = self.act(x) |
| | x = self.c_proj(x) |
| | return self.dropout(x) |
| | |
| | class GPT2Block(nn.Module): |
| | def __init__(self,config:SimpleNamespace): |
| | super(GPT2Block,self).__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:torch.Tensor,enc_out:torch.Tensor)->torch.Tensor: |
| | 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:SimpleNamespace): |
| | super(VisionGPT2Model,self).__init__() |
| | self.config = config |
| | 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.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:torch.Tensor)->torch.Tensor: |
| | 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 x |
| | |
| | def pretrained_layers_trainable(self,t:bool = False)->None: |
| | 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 = t |
| | else: |
| | layer.requires_grad = t |
| | |
| | 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)->None: |
| | 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:SimpleNamespace): |
| | 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:torch.Tensor,input_ids:torch.Tensor,labels:None|torch.Tensor=None)->torch.Tensor: |
| | image = self.patch_embed(image) |
| | image = self._pos_embed(image) |
| | token_embeddings = self.transformer.wte(input_ids) |
| | pos_embs = torch.arange(0,input_ids.size(1)).to(input_ids.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:torch.Tensor, |
| | sequence:torch.Tensor, |
| | max_tokens:int =50, |
| | temp:float =1.0, |
| | deter:bool =False) -> torch.Tensor: |
| | |
| | for _ in range(max_tokens): |
| | out = self(image,sequence) |
| | out = out[:,-1,:]/temp |
| | probs = F.softmax(out,dim=-1) |
| | if deter: |
| | 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() |
