Skip to content
Snippets Groups Projects
Commit c841c22a authored by kulcsar's avatar kulcsar
Browse files

add changes to main and train for tmix

parent 0eb1bf33
No related branches found
No related tags found
No related merge requests found
......@@ -38,7 +38,7 @@ def run(raw_args):
else:
print("non eligible model type selected")
elif args.model_type == "one":
model=models.WordClassificationModel(args.architecture).to("cuda")
model=models.WordClassificationModel(args.architecture, args.tmix, args.mixlayer, args.lambda_value).to("cuda")
else:
print("non eligible model type selected")
......@@ -67,7 +67,7 @@ def run(raw_args):
#train...
print("training..")
if args.train_loop=="swp":
evaluation_test, evaluation_train = train.train(model, args.architecture,args.random_seed, args.gradient_accumulation_steps, args.mix_up, args.threshold, args.lambda_value, args.mixup_epoch, train_dataset, test_dataset, args.epochs, args.learning_rate, args.batch_size, args.test_batch_size)
evaluation_test, evaluation_train = train.train(model, args.architecture,args.random_seed, args.gradient_accumulation_steps, args.mix_up, args.threshold, args.lambda_value, args.mixup_epoch, args.tmix, args.mixlayer, train_dataset, test_dataset, args.epochs, args.learning_rate, args.batch_size, args.test_batch_size)
elif args.train_loop=="salami":
evaluation_test = train.train_salami(model,args.random_seed, train_dataset, test_dataset, args.batch_size, args.test_batch_size, args.learning_rate, args.epochs)
else:
......@@ -93,6 +93,18 @@ if __name__ == "__main__":
"--model_type",
help="How to initialize the Classification Model",
choices=["separate", "one"])
parser.add_argument(
"--tmix",
help="whether or not to use tmix. if yes, please specify layer and lambda"
action="store-true"
)
parser.add_argument(
"--mixlayer",
help="specify the layer to mix. Only select one layer at a time"
type=list
)
#Datasets
parser.add_argument(
......
......@@ -9,6 +9,8 @@ import math
from tqdm.auto import tqdm
from transformers import BertTokenizer, RobertaTokenizer, BertModel, RobertaModel, RobertaPreTrainedModel, RobertaConfig, BertConfig, BertPreTrainedModel, PreTrainedModel, AutoModel, AutoTokenizer, AutoConfig
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, BaseModelOutputWithPoolingAndCrossAttentions, CausalLMOutputWithCrossAttentions, MaskedLMOutput, MultipleChoiceModelOutput, NextSentencePredictorOutput, QuestionAnsweringModelOutput, SequenceClassifierOutput, TokenClassifierOutput
from transformers.models.bert.modeling_bert import BertEmbeddings, BertLayer, BertPooler
from transformers import AdamW, get_scheduler
from torch import nn
from torch.nn import CrossEntropyLoss
......@@ -16,6 +18,7 @@ import matplotlib.pyplot as plt
import os
import pandas as pd
import sklearn
from typing import List, Optional, Tuple, Union
metric=evaluate.load("accuracy")
torch.cuda.empty_cache()
......@@ -38,11 +41,14 @@ class WordClassificationModel(torch.nn.Module): #AutoModel verwenden aus der Bib
and the linear classifier layer to make it a binary decision problem. In the forward step
we specify the classification over the span given by end and start position and compute the
loss function with cross entropy. The predictions (logits) are made by our classifier layer."""
def __init__(self, config_name):
def __init__(self, config_name, tmix=False, mixlayer=-1, lambda_value=0.0):
super(WordClassificationModel, self).__init__()
#self.num_labels=config.num_labels
self.embedding_model=AutoModel.from_pretrained(config_name, config=AutoConfig.from_pretrained(config_name))
if tmix:
print("initializing BertModelTMix")
self.embedding_model=BertModelTMix(config=AutoConfig.from_pretrained(config_name), mixlayer=mixlayer, lambda_value=lambda_value)
else:
self.embedding_model=AutoModel.from_pretrained(config_name, config=AutoConfig.from_pretrained(config_name))
self.dropout=nn.Dropout(0.1) #config.hidden_dropout_prob for BERT: 0.1, config.hidden_dropout_prob:Roberta: not in roberta
......@@ -159,3 +165,340 @@ class RobertaForWordClassification(RobertaPreTrainedModel): #AutoModel verwenden
return outputs
class BertModelTMix(BertPreTrainedModel):
"""
The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
cross-attention is added between the self-attention layers, following the architecture described in [Attention is
all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
`add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
"""
def __init__(self, config, add_pooling_layer=True, lambda_value=0.0, mixlayer=-1):
super().__init__(config)
self.config = config
self.embeddings = BertEmbeddings(config)
self.encoder = BertTMixEncoder(config, lambda_value, mixlayer)
self.pooler = BertPooler(config) if add_pooling_layer else None
# Initialize weights and apply final processing
self.post_init()
def get_input_embeddings(self):
return self.embeddings.word_embeddings
def set_input_embeddings(self, value):
self.embeddings.word_embeddings = value
def _prune_heads(self, heads_to_prune):
"""
Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
class PreTrainedModel
"""
for layer, heads in heads_to_prune.items():
self.encoder.layer[layer].attention.prune_heads(heads)
def post_init(self):
print("jadijaid")
#@add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
#@add_code_sample_docstrings(
# checkpoint=_CHECKPOINT_FOR_DOC,
# output_type=BaseModelOutputWithPoolingAndCrossAttentions,
# config_class=_CONFIG_FOR_DOC,
#)
def forward(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
token_type_ids: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
encoder_hidden_states: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
r"""
encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
the model is configured as a decoder.
encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**.
past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
`decoder_input_ids` of shape `(batch_size, sequence_length)`.
use_cache (`bool`, *optional*):
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
`past_key_values`).
"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if self.config.is_decoder:
use_cache = use_cache if use_cache is not None else self.config.use_cache
else:
use_cache = False
if input_ids is not None and inputs_embeds is not None:
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
elif input_ids is not None:
input_shape = input_ids.size()
elif inputs_embeds is not None:
input_shape = inputs_embeds.size()[:-1]
else:
raise ValueError("You have to specify either input_ids or inputs_embeds")
batch_size, seq_length = input_shape
device = input_ids.device if input_ids is not None else inputs_embeds.device
# past_key_values_length
past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
if attention_mask is None:
attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
if token_type_ids is None:
if hasattr(self.embeddings, "token_type_ids"):
buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
token_type_ids = buffered_token_type_ids_expanded
else:
token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
# ourselves in which case we just need to make it broadcastable to all heads.
extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device=None)
# If a 2D or 3D attention mask is provided for the cross-attention
# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
if self.config.is_decoder and encoder_hidden_states is not None:
encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
if encoder_attention_mask is None:
encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
else:
encoder_extended_attention_mask = None
# Prepare head mask if needed
# 1.0 in head_mask indicate we keep the head
# attention_probs has shape bsz x n_heads x N x N
# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
embedding_output = self.embeddings(
input_ids=input_ids,
position_ids=position_ids,
token_type_ids=token_type_ids,
inputs_embeds=inputs_embeds,
past_key_values_length=past_key_values_length,
)
encoder_outputs = self.encoder(
embedding_output,
attention_mask=extended_attention_mask,
head_mask=head_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_extended_attention_mask,
past_key_values=past_key_values,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
sequence_output = encoder_outputs[0]
pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
if not return_dict:
return (sequence_output, pooled_output) + encoder_outputs[1:]
return BaseModelOutputWithPoolingAndCrossAttentions(
last_hidden_state=sequence_output,
pooler_output=pooled_output,
past_key_values=encoder_outputs.past_key_values,
hidden_states=encoder_outputs.hidden_states,
attentions=encoder_outputs.attentions,
cross_attentions=encoder_outputs.cross_attentions,
)
class BertTMixEncoder(torch.nn.Module):
"""Used for Tmix. When using Tmix the only change that has to be done, is to be able to modify at which layer to start
training and what the input of that layer is. We can do so, by making i in forward a variable. However, how do we specify
the input of the layer?"""
def __init__(self, config, lambda_value, mixlayer):
super().__init__()
self.config = config
BertLayer.forward=forward_new(BertLayer.forward) #Monkey Patch
self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
self.gradient_checkpointing = False
self.lambda_value=lambda_value
self.mixlayer=mixlayer
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = False,
output_hidden_states: Optional[bool] = False,
return_dict: Optional[bool] = True,
start_layer: Optional[int] = 0) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
all_hidden_states = () if output_hidden_states else None
all_self_attentions = () if output_attentions else None
all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
#self.layer.forward=forward_new(self.layer.forward) #Monkeypatch here?
next_decoder_cache = () if use_cache else None
for i, layer_module in enumerate(self.layer):
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
layer_head_mask = head_mask[i] if head_mask is not None else None
past_key_value = past_key_values[i] if past_key_values is not None else None
if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
)
use_cache = False
def create_custom_forward(module):
def custom_forward(*inputs):
return module(*inputs, past_key_value, output_attentions)
return custom_forward
layer_outputs = torch.utils.checkpoint.checkpoint(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
layer_head_mask,
encoder_hidden_states,
encoder_attention_mask,
)
else:
print("before")
layer_outputs = layer_module(
hidden_states,
attention_mask,
layer_head_mask,
encoder_hidden_states,
encoder_attention_mask,
past_key_value,
output_attentions,
lambda_value=self.lambda_value,
mixlayer=self.mixlayer,
layer_now=i
)
print("after")
hidden_states = layer_outputs[0]
print(hidden_states.size())
if use_cache:
next_decoder_cache += (layer_outputs[-1],)
if output_attentions:
all_self_attentions = all_self_attentions + (layer_outputs[1],)
if self.config.add_cross_attention:
all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
if not return_dict:
return tuple(
v
for v in [
hidden_states,
next_decoder_cache,
all_hidden_states,
all_self_attentions,
all_cross_attentions,
]
if v is not None
)
return BaseModelOutputWithPastAndCrossAttentions(
last_hidden_state=hidden_states,
past_key_values=next_decoder_cache,
hidden_states=all_hidden_states,
attentions=all_self_attentions,
cross_attentions=all_cross_attentions,
)
#Moneky Patching the forward function of BertLayer for mixup -> use decorators here to call the old forward function on the newly comptued hidden_state
def forward_new(forward):
print("over here")
def forward_mix(self, hidden_states: torch.Tensor,
attention_mask: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = False,
output_hidden_states: Optional[bool] = False,
return_dict: Optional[bool] = True,
lambda_value: float=0.0,
mixlayer: int=-1,
layer_now: int=0)-> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
print("in Monkey Patch function")
#print("attention mask:", attention_mask.size())
#print("hidden states: ", hidden_states.size())
#print("past key values: ", past_key_values.size())
#print("head mask: ", head_mask.size())
#print("encoder hidden states: ", encoder_hidden_states.size())
if layer_now == mixlayer
runs = math.floor(hidden_states.size()[0]/2)
print("runs: ", runs)
print("lambda_value: ", lambda_value)
counter=0
new_matrices=[]
new_attention_masks=[]
for i in range(runs):
hidden_states_1=hidden_states[counter]
attention_mask_1=attention_mask[counter]
hidden_states_2=hidden_states[counter+1]
attention_mask_2=attention_mask[counter+1]
new_matrix = (lambda_value*hidden_states_1) + ((1-lambda_value)*hidden_states_2) #do interpolation
new_attention_mask = (lambda_value*attention_mask*attention_mask_1) + ((1-lambda_value)*attention_mask_2)
new_attention_masks.append(new_attention_mask)
new_matrices.append(new_matrix) #do intepolations
for i in range(hidden_states.size()[0]-runs):
new_matrices.append(torch.zeros([512, 768]))
new_matrices=torch.stack(new_matrices)
new_attention_masks=torch.stack(new_attention_masks)
print(new_matrices.size())
outputs=forward(self, hidden_states=new_matrices, head_mask=head_mask, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask, past_key_value=past_key_values, output_attentions=output_attentions) #I"m a bit confused here... do we have to add self or rather not?
else:
outputs = forward(self, hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, past_key_value=past_key_values, output_attentions)
print(outputs)
return outputs
return forward_mix
......@@ -22,7 +22,7 @@ torch.cuda.empty_cache()
#with torch.autocast("cuda"):
def train(model, name, seed,gradient_accumulation_steps,mixup, threshold, lambda_value,mixup_epoch, train_dataset, test_dataset, num_epochs, learning_rate, batch_size, test_batch_size):
def train(model, name, seed,gradient_accumulation_steps,mixup, threshold, lambda_value, mixup_epoch, tmix, mixlayer, train_dataset, test_dataset, num_epochs, learning_rate, batch_size, test_batch_size):
"""Write Train loop for model with certain train dataset"""
#set_seed(seed)
#if model_name[0] == "b":
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment