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()