Repository: flairNLP/flair

Issue:
Change SequenceTagger and DocumentRNNEmbeddings to batch_first mode
We are currently not consistent in Flair with whether we use `batch_first=True` or not. The `LockedDropout` and `WordDropout` classes for instance do not currently support a `batch_first` mode, so we do awkward transpose operations in the `SequenceTagger` and `DocumentRNNEmbeddings`. 

**Task**: Make everything `batch_first`. Add support for this in `LockedDropout` and `WordDropout` and remove transpose operations from `SequenceTagger` and `DocumentRNNEmbeddings`.



Discussion:

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -2681,6 +2681,7 @@ def __init__(
                 hidden_size,
                 num_layers=rnn_layers,
                 bidirectional=self.bidirectional,
+                batch_first=True,
             )
         else:
             self.rnn = torch.nn.GRU(
@@ -2688,6 +2689,7 @@ def __init__(
                 hidden_size,
                 num_layers=rnn_layers,
                 bidirectional=self.bidirectional,
+                batch_first=True,
             )
 
         self.name = "document_" + self.rnn._get_name()
@@ -2721,21 +2723,11 @@ def _add_embeddings_internal(self, sentences: List[Sentence]):
 
         self.rnn.zero_grad()
 
-        # the permutation that sorts the sentences by length, descending
-        sort_perm = np.argsort([len(s) for s in sentences])[::-1]
-
-        # the inverse permutation that restores the input order; it's an index tensor therefore LongTensor
-        sort_invperm = np.argsort(sort_perm)
-
-        # sort sentences by number of tokens
-        sentences = [sentences[i] for i in sort_perm]
-
+        # embed words in the sentence
         self.embeddings.embed(sentences)
 
-        longest_token_sequence_in_batch: int = len(sentences[0])
-
-        # all_sentence_tensors = []
-        lengths: List[int] = []
+        lengths: List[int] = [len(sentence.tokens) for sentence in sentences]
+        longest_token_sequence_in_batch: int = max(lengths)
 
         # initialize zero-padded word embeddings tensor
         sentence_tensor = torch.zeros(
@@ -2748,20 +2740,16 @@ def _add_embeddings_internal(self, sentences: List[Sentence]):
             device=flair.device,
         )
 
-        # fill values with word embeddings
         for s_id, sentence in enumerate(sentences):
-            lengths.append(len(sentence.tokens))
-
+            # fill values with word embeddings
             sentence_tensor[s_id][: len(sentence)] = torch.cat(
                 [token.get_embedding().unsqueeze(0) for token in sentence], 0
             )
 
-        # TODO: this can only be removed once the implementations of word_dropout and locked_dropout have a batch_first mode
-        sentence_tensor = sentence_tensor.transpose(0, 1)
-
         # --------------------------------------------------------------------
         # FF PART
         # --------------------------------------------------------------------
+        sentence_tensor = self.dropout(sentence_tensor)
         # use word dropout if set
         if self.use_word_dropout:
             sentence_tensor = self.word_dropout(sentence_tensor)
@@ -2769,14 +2757,13 @@ def _add_embeddings_internal(self, sentences: List[Sentence]):
         if self.reproject_words:
             sentence_tensor = self.word_reprojection_map(sentence_tensor)
 
-        sentence_tensor = self.dropout(sentence_tensor)
-        packed = pack_padded_sequence(sentence_tensor, lengths)
-
-        self.rnn.flatten_parameters()
+        packed = pack_padded_sequence(
+            sentence_tensor, lengths, enforce_sorted=False, batch_first=True
+        )
 
         rnn_out, hidden = self.rnn(packed)
 
-        outputs, output_lengths = pad_packed_sequence(rnn_out)
+        outputs, output_lengths = pad_packed_sequence(rnn_out, batch_first=True)
 
         outputs = self.dropout(outputs)
 
@@ -2784,11 +2771,11 @@ def _add_embeddings_internal(self, sentences: List[Sentence]):
         # EXTRACT EMBEDDINGS FROM RNN
         # --------------------------------------------------------------------
         for sentence_no, length in enumerate(lengths):
-            last_rep = outputs[length - 1, sentence_no]
+            last_rep = outputs[sentence_no, length - 1]
 
             embedding = last_rep
             if self.bidirectional:
-                first_rep = outputs[0, sentence_no]
+                first_rep = outputs[sentence_no, 0]
                 embedding = torch.cat([first_rep, last_rep], 0)
 
             if self.static_embeddings:
@@ -2797,9 +2784,6 @@ def _add_embeddings_internal(self, sentences: List[Sentence]):
             sentence = sentences[sentence_no]
             sentence.set_embedding(self.name, embedding)
 
-        # restore original order of sentences in the batch
-        sentences = [sentences[i] for i in sort_invperm]
-
 
 @deprecated(
     version="0.4",
diff --git a/flair/models/sequence_tagger_model.py b/flair/models/sequence_tagger_model.py
--- a/flair/models/sequence_tagger_model.py
+++ b/flair/models/sequence_tagger_model.py
@@ -154,6 +154,7 @@ def __init__(
                     num_layers=self.nlayers,
                     dropout=0.0 if self.nlayers == 1 else 0.5,
                     bidirectional=True,
+                    batch_first=True,
                 )
                 # Create initial hidden state and initialize it
                 if self.train_initial_hidden_state:
@@ -454,9 +455,6 @@ def forward(self, sentences: List[Sentence]):
                 [token.get_embedding().unsqueeze(0) for token in sentence], 0
             )
 
-        # TODO: this can only be removed once the implementations of word_dropout and locked_dropout have a batch_first mode
-        sentence_tensor = sentence_tensor.transpose(0, 1)
-
         # --------------------------------------------------------------------
         # FF PART
         # --------------------------------------------------------------------
@@ -472,7 +470,7 @@ def forward(self, sentences: List[Sentence]):
 
         if self.use_rnn:
             packed = torch.nn.utils.rnn.pack_padded_sequence(
-                sentence_tensor, lengths, enforce_sorted=False
+                sentence_tensor, lengths, enforce_sorted=False, batch_first=True
             )
 
             # if initial hidden state is trainable, use this state
@@ -496,9 +494,6 @@ def forward(self, sentences: List[Sentence]):
             #     sentence_tensor = self.word_dropout(sentence_tensor)
             if self.use_locked_dropout > 0.0:
                 sentence_tensor = self.locked_dropout(sentence_tensor)
-        else:
-            # transpose to batch_first mode
-            sentence_tensor = sentence_tensor.transpose(0, 1)
 
         features = self.linear(sentence_tensor)
 
diff --git a/flair/nn.py b/flair/nn.py
--- a/flair/nn.py
+++ b/flair/nn.py
@@ -144,16 +144,21 @@ class LockedDropout(torch.nn.Module):
     Implementation of locked (or variational) dropout. Randomly drops out entire parameters in embedding space.
     """
 
-    def __init__(self, dropout_rate=0.5, inplace=False):
+    def __init__(self, dropout_rate=0.5, batch_first=True, inplace=False):
         super(LockedDropout, self).__init__()
         self.dropout_rate = dropout_rate
+        self.batch_first = batch_first
         self.inplace = inplace
 
     def forward(self, x):
         if not self.training or not self.dropout_rate:
             return x
 
-        m = x.data.new(1, x.size(1), x.size(2)).bernoulli_(1 - self.dropout_rate)
+        if not self.batch_first:
+            m = x.data.new(1, x.size(1), x.size(2)).bernoulli_(1 - self.dropout_rate)
+        else:
+            m = x.data.new(x.size(0), 1, x.size(2)).bernoulli_(1 - self.dropout_rate)
+
         mask = torch.autograd.Variable(m, requires_grad=False) / (1 - self.dropout_rate)
         mask = mask.expand_as(x)
         return mask * x
@@ -177,9 +182,9 @@ def forward(self, x):
         if not self.training or not self.dropout_rate:
             return x
 
-        m = x.data.new(x.size(0), 1, 1).bernoulli_(1 - self.dropout_rate)
+        m = x.data.new(x.size(0), x.size(1), 1).bernoulli_(1 - self.dropout_rate)
+
         mask = torch.autograd.Variable(m, requires_grad=False)
-        mask = mask.expand_as(x)
         return mask * x
 
     def extra_repr(self):


Test patch: