Optimized wrapper with correct API

modeling_norbert.py

@@ -1,12 +1,9 @@
-from __future__ import absolute_import, division, print_function, unicode_literals
-
 import math
 from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch import _softmax_backward_data as _softmax_backward_data
 from torch.utils import checkpoint
 
 from configuration_norbert import NorbertConfig
@@ -20,6 +17,7 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
     BaseModelOutput
 )
+from transformers.pytorch_utils import softmax_backward_data
 
 
 class Encoder(nn.Module):
@@ -130,8 +128,8 @@ class MaskedSoftmax(torch.autograd.Function):
     @staticmethod
     def backward(self, grad_output):
         output, = self.saved_tensors
-        inputGrad = _softmax_backward_data(grad_output, output, self.dim, output.dtype)
-        return inputGrad, None, None
+        input_grad = softmax_backward_data(self, grad_output, output, self.dim, output)
+        return input_grad, None, None
 
 
 class Attention(nn.Module):
@@ -188,31 +186,36 @@ class Attention(nn.Module):
         if self.position_indices.size(0) < query_len:
             position_indices = torch.arange(query_len, dtype=torch.long).unsqueeze(1) \
                 - torch.arange(query_len, dtype=torch.long).unsqueeze(0)
-            position_indices = self.make_log_bucket_position(position_indices, self.config.position_bucket_size, 512)
-            position_indices = self.config.position_bucket_size - 1 + position_indices
-            self.register_buffer("position_indices", position_indices.to(hidden_states.device), persistent=True)
+            position_indices = self.make_log_bucket_position(position_indices, self.position_bucket_size, 512)
+            position_indices = self.position_bucket_size - 1 + position_indices
+            self.position_indices = position_indices.to(hidden_states.device)
 
         hidden_states = self.pre_layer_norm(hidden_states)
 
         query, key = self.in_proj_qk(hidden_states).chunk(2, dim=2)  # shape: [T, B, D]
         value = self.in_proj_v(hidden_states)  # shape: [T, B, D]
 
-        pos = self.in_proj_qk(self.dropout(relative_embedding))  # shape: [2T-1, 2D]
-        pos = F.embedding(self.position_indices[:query_len, :key_len], pos)  # shape: [T, T, 2D]
-        pos = pos.view(query_len, key_len, self.num_heads, 2*self.head_size)
-        query_pos, key_pos = pos.chunk(2, dim=3)
-
         query = query.reshape(query_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
         key = key.reshape(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
         value = value.view(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
 
         attention_scores = torch.bmm(query, key.transpose(1, 2) * self.scale)
 
+        pos = self.in_proj_qk(self.dropout(relative_embedding))  # shape: [2T-1, 2D]
+        query_pos, key_pos = pos.view(-1, self.num_heads, 2*self.head_size).chunk(2, dim=2)
         query = query.view(batch_size, self.num_heads, query_len, self.head_size)
         key = key.view(batch_size, self.num_heads, query_len, self.head_size)
+
+        attention_c_p = torch.einsum("bhqd,khd->bhqk", query, key_pos.squeeze(1) * self.scale)
+        attention_p_c = torch.einsum("bhkd,qhd->bhqk", key * self.scale, query_pos.squeeze(1))
+
+        position_indices = self.position_indices[:query_len, :key_len].expand(batch_size, self.num_heads, -1, -1)
+        attention_c_p = attention_c_p.gather(3, position_indices)
+        attention_p_c = attention_p_c.gather(2, position_indices)
+
         attention_scores = attention_scores.view(batch_size, self.num_heads, query_len, key_len)
-        attention_scores.add_(torch.einsum("bhqd,qkhd->bhqk", query, key_pos * self.scale))
-        attention_scores.add_(torch.einsum("bhkd,qkhd->bhqk", key * self.scale, query_pos))
+        attention_scores.add_(attention_c_p)
+        attention_scores.add_(attention_p_c)
 
         return attention_scores, value
 
@@ -332,12 +335,16 @@ class NorbertModel(NorbertPreTrainedModel):
         sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
 
         if not return_dict:
-            return sequence_output, contextualized_embeddings, attention_probs
+            return (
+                sequence_output,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else [])
+            )
 
         return BaseModelOutput(
             last_hidden_state=sequence_output,
-            hidden_states=contextualized_embeddings,
-            attentions=attention_probs
+            hidden_states=contextualized_embeddings if output_hidden_states else None,
+            attentions=attention_probs if output_attentions else None
         )
 
 
@@ -375,14 +382,18 @@ class NorbertForMaskedLM(NorbertModel):
             masked_lm_loss = F.cross_entropy(subword_prediction.flatten(0, 1), labels.flatten())
 
         if not return_dict:
-            output = (subword_prediction, contextualized_embeddings, attention_probs)
+            output = (
+                subword_prediction,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else [])
+            )
             return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
 
         return MaskedLMOutput(
             loss=masked_lm_loss,
             logits=subword_prediction,
-            hidden_states=contextualized_embeddings,
-            attentions=attention_probs
+            hidden_states=contextualized_embeddings if output_hidden_states else None,
+            attentions=attention_probs if output_attentions else None
         )
 
 
@@ -465,14 +476,18 @@ class NorbertForSequenceClassification(NorbertModel):
                 loss = loss_fct(logits, labels)
 
         if not return_dict:
-            output = (logits, contextualized_embeddings, attention_probs)
+            output = (
+                logits,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else [])
+            )
             return ((loss,) + output) if loss is not None else output
 
         return SequenceClassifierOutput(
             loss=loss,
             logits=logits,
-            hidden_states=contextualized_embeddings,
-            attentions=attention_probs
+            hidden_states=contextualized_embeddings if output_hidden_states else None,
+            attentions=attention_probs if output_attentions else None
         )
 
 
@@ -508,14 +523,18 @@ class NorbertForTokenClassification(NorbertModel):
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
 
         if not return_dict:
-            output = (logits, contextualized_embeddings, attention_probs)
+            output = (
+                logits,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else [])
+            )
             return ((loss,) + output) if loss is not None else output
 
         return TokenClassifierOutput(
             loss=loss,
             logits=logits,
-            hidden_states=contextualized_embeddings,
-            attentions=attention_probs
+            hidden_states=contextualized_embeddings if output_hidden_states else None,
+            attentions=attention_probs if output_attentions else None
         )
 
 
@@ -569,15 +588,20 @@ class NorbertForQuestionAnswering(NorbertModel):
             total_loss = (start_loss + end_loss) / 2
 
         if not return_dict:
-            output = start_logits, end_logits, contextualized_embeddings, attention_probs
+            output = (
+                start_logits,
+                end_logits,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else [])
+            )
             return ((total_loss,) + output) if total_loss is not None else output
 
         return QuestionAnsweringModelOutput(
             loss=total_loss,
             start_logits=start_logits,
             end_logits=end_logits,
-            hidden_states=contextualized_embeddings,
-            attentions=attention_probs,
+            hidden_states=contextualized_embeddings if output_hidden_states else None,
+            attentions=attention_probs if output_attentions else None
         )
 
 
@@ -598,9 +622,9 @@ class NorbertForMultipleChoice(NorbertModel):
         token_type_ids: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         labels: Optional[torch.Tensor] = None,
-        return_dict: Optional[bool] = None,
-        start_positions: Optional[torch.Tensor] = None,
-        end_positions: Optional[torch.Tensor] = None
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None
     ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         num_choices = input_ids.shape[1]
@@ -618,12 +642,16 @@ class NorbertForMultipleChoice(NorbertModel):
             loss = loss_fct(reshaped_logits, labels)
 
         if not return_dict:
-            output = (reshaped_logits, contextualized_embeddings, attention_probs)
+            output = (
+                reshaped_logits,
+                *([contextualized_embeddings] if output_hidden_states else []),
+                *([attention_probs] if output_attentions else [])
+            )
             return ((loss,) + output) if loss is not None else output
 
         return MultipleChoiceModelOutput(
             loss=loss,
             logits=reshaped_logits,
-            hidden_states=contextualized_embeddings,
-            attentions=attention_probs,
+            hidden_states=contextualized_embeddings if output_hidden_states else None,
+            attentions=attention_probs if output_attentions else None
         )