add changes to main and train for tmix

Code/main.py

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

Code/models.py

@@ -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

Code/train.py

@@ -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":