Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from torch.nn import CrossEntropyLoss | |
| from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2DoubleHeadsModel | |
| from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions | |
| from types import MethodType | |
| tokenizer = GPT2Tokenizer.from_pretrained('gpt2') | |
| def get_model(device='cpu', gpt2_name='gpt2', in_net=False, in_net_init_identity=True, out_net=False, out_net_init_identity=True, freeze_ln=False, freeze_pos=True, | |
| freeze_wte=True, freeze_ff=True, freeze_attn=True, dup_lm_head=False, dup_lm_head_bias=False): | |
| # ['gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'] | |
| model = GPT2LMHeadModel.from_pretrained(gpt2_name).to(device) | |
| # model = GPT2DoubleHeadsModel.from_pretrained('gpt2') | |
| """ | |
| Initialize linear input layer | |
| """ | |
| in_layer_sizes = [] | |
| out_layer_sizes = [] | |
| input_dim = model.config.n_embd | |
| dropout = 0.1 | |
| orth_gain = 1.41 | |
| # orth_gain = None | |
| in_net_init_identity = True | |
| #Model - in_net | |
| if in_net: | |
| in_layers = [] | |
| last_output_size = input_dim | |
| for size in in_layer_sizes: | |
| layer = nn.Linear(last_output_size, size) | |
| if orth_gain is not None: | |
| torch.nn.init.orthogonal_(layer.weight, gain=orth_gain) | |
| layer.bias.data.zero_() | |
| in_layers.append(layer) | |
| in_layers.append(nn.ReLU()) | |
| in_layers.append(nn.Dropout(dropout)) | |
| last_output_size = size | |
| in_final_linear = nn.Linear(last_output_size, model.config.n_embd) | |
| # if orth_gain is not None: | |
| # torch.nn.init.orthogonal_(in_final_linear.weight, gain=orth_gain) | |
| # in_final_linear.bias.data.zero_() | |
| # Initialize final_linear layer to identity transformation | |
| if in_net_init_identity: | |
| nn.init.eye_(in_final_linear.weight) | |
| in_final_linear.bias.data.zero_() | |
| in_layers.append(in_final_linear) | |
| in_layers.append(nn.Dropout(dropout)) | |
| model.in_net = nn.Sequential(*in_layers) | |
| model.in_net.requires_grad = True | |
| """ | |
| Initialize linear output layer | |
| """ | |
| if out_net: | |
| output_dim = model.config.n_embd | |
| out_layers = [] | |
| last_output_size = model.config.n_embd | |
| for size in out_layer_sizes: | |
| out_layers.append(nn.Linear(last_output_size, size)) | |
| out_layers.append(nn.ReLU()) | |
| out_layers.append(nn.Dropout(dropout)) | |
| last_output_size = size | |
| out_final_linear = nn.Linear(last_output_size, output_dim) | |
| if out_net_init_identity: | |
| nn.init.eye_(out_final_linear.weight) | |
| out_final_linear.bias.data.zero_() | |
| out_layers.append(out_final_linear) | |
| model.out_net = nn.Sequential(*out_layers) | |
| model.out_net.requires_grad = True | |
| """ | |
| out layer on top of lm_head | |
| """ | |
| # out_net_top = nn.Linear(model.config.vocab_size, model.config.vocab_size) | |
| # nn.init.eye_(out_net_top.weight) | |
| # model.out_net_top = out_net_top | |
| # model.out_net_top.requires_grad = True | |
| if dup_lm_head: | |
| lm_head_new = nn.Linear(model.config.n_embd, | |
| model.config.vocab_size, bias=dup_lm_head_bias) | |
| lm_head_new.weight = torch.nn.Parameter( | |
| model.lm_head.weight.data.detach().clone(), requires_grad=True) | |
| # lm_head_new.bias.data.zero_() | |
| model.lm_head_new = lm_head_new | |
| model.lm_head_new.requires_grad = True | |
| """ | |
| Freeze transformer layers | |
| """ | |
| total_parameters = 0 | |
| target_parameters = 0 | |
| for name, p in model.transformer.named_parameters(): | |
| name = name.lower() | |
| size = p.size() | |
| param_count = 1 | |
| for dimension in size: | |
| param_count *= dimension | |
| total_parameters += param_count | |
| if 'ln' in name or 'norm' in name: | |
| p.requires_grad = not freeze_ln | |
| elif 'wpe' in name or 'position_embeddings' in name or 'pos_drop' in name: | |
| p.requires_grad = not freeze_pos | |
| target_parameters += param_count | |
| elif 'mlp' in name: | |
| p.requires_grad = not freeze_ff | |
| elif 'attn' in name: | |
| p.requires_grad = not freeze_attn | |
| elif 'wte' in name: | |
| p.requires_grad = not freeze_wte | |
| else: | |
| p.requires_grad = False | |
| # print(f'Total params: {total_parameters}') | |
| # print( | |
| # f'Target params: {target_parameters} ({target_parameters / total_parameters * 100:.2f}%)') | |
| def forward( | |
| self, | |
| input_ids=None, | |
| past_key_values=None, | |
| attention_mask=None, | |
| token_type_ids=None, | |
| position_ids=None, | |
| head_mask=None, | |
| inputs_embeds=None, | |
| encoder_hidden_states=None, | |
| encoder_attention_mask=None, | |
| labels=None, | |
| use_cache=None, | |
| output_attentions=None, | |
| output_hidden_states=None, | |
| return_dict=None, | |
| **kwargs | |
| ): | |
| r""" | |
| labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): | |
| Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set | |
| ``labels = input_ids`` Indices are selected in ``[-100, 0, ..., config.vocab_size]`` All labels set to | |
| ``-100`` are ignored (masked), the loss is only computed for labels in ``[0, ..., config.vocab_size]`` | |
| """ | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
| # Convert from input ids to word embeddings so that we can apply a linear layer | |
| x = self.transformer.wte(input_ids) | |
| try: | |
| x = self.in_net(x) | |
| except AttributeError: | |
| pass | |
| transformer_outputs = self.transformer( | |
| inputs_embeds=x, | |
| past_key_values=past_key_values, | |
| attention_mask=attention_mask, | |
| token_type_ids=token_type_ids, | |
| position_ids=position_ids, | |
| head_mask=head_mask, | |
| encoder_hidden_states=encoder_hidden_states, | |
| encoder_attention_mask=encoder_attention_mask, | |
| use_cache=use_cache, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict, | |
| **kwargs | |
| ) | |
| hidden_states = transformer_outputs[0] | |
| # Set device for model parallelism | |
| if self.model_parallel: | |
| torch.cuda.set_device(self.transformer.first_device) | |
| hidden_states = hidden_states.to(self.lm_head.weight.device) | |
| try: | |
| hidden_states = self.out_net(hidden_states) | |
| except AttributeError: | |
| pass | |
| try: | |
| lm_logits = self.lm_head_new(hidden_states) | |
| except AttributeError: | |
| lm_logits = self.lm_head(hidden_states) | |
| # lm_logits = self.out_net_top(lm_logits) | |
| loss = None | |
| if labels is not None: | |
| # Shift so that tokens < n predict n | |
| shift_logits = lm_logits[..., :-1, :].contiguous() | |
| shift_labels = labels[..., 1:].contiguous() | |
| # Flatten the tokens | |
| loss_fct = CrossEntropyLoss() | |
| loss = loss_fct( | |
| shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) | |
| if not return_dict: | |
| output = (lm_logits,) + transformer_outputs[1:] | |
| return ((loss,) + output) if loss is not None else output | |
| return CausalLMOutputWithCrossAttentions( | |
| loss=loss, | |
| logits=lm_logits, | |
| past_key_values=transformer_outputs.past_key_values, | |
| hidden_states=transformer_outputs.hidden_states, | |
| attentions=transformer_outputs.attentions, | |
| cross_attentions=transformer_outputs.cross_attentions, | |
| ) | |
| model.forward = MethodType(forward, model) | |
| return model | |
| # model = get_model() | |
| ''' | |
| only for testing purpose | |
| ''' | |
| if __name__ == "__main__": | |
| model = get_model(gpt2_name='gpt2', in_net=False, in_net_init_identity=True, out_net=False, out_net_init_identity=False, freeze_ln=True, freeze_pos=True, | |
| freeze_wte=True, freeze_ff=True, freeze_attn=True) | |
| for name, p in model.named_parameters(): | |
| if p.requires_grad: | |
| print(name, p.requires_grad) | |
| for p in model.lm_head_new.parameters(): | |
| print('lm_head_new', p) | |
| # for p in model.out_net.parameters(): | |
| # print('out_net',p) | |