Optimize AAN transformer and small fixes (#1482)

* Optimize AAN transformer and small fixes
* Make use of FFN layer in AAN an option

Author: vince62s

onmt/decoders/transformer.py

@@ -23,15 +23,17 @@ class TransformerDecoderLayer(nn.Module):
     """
 
     def __init__(self, d_model, heads, d_ff, dropout,
-                 self_attn_type="scaled-dot", max_relative_positions=0):
+                 self_attn_type="scaled-dot", max_relative_positions=0,
+                 aan_useffn=False):
         super(TransformerDecoderLayer, self).__init__()
 
         if self_attn_type == "scaled-dot":
             self.self_attn = MultiHeadedAttention(
                 heads, d_model, dropout=dropout,
                 max_relative_positions=max_relative_positions)
         elif self_attn_type == "average":
-            self.self_attn = AverageAttention(d_model, dropout=dropout)
+            self.self_attn = AverageAttention(d_model, dropout=dropout,
+                                              aan_useffn=aan_useffn)
 
         self.context_attn = MultiHeadedAttention(
             heads, d_model, dropout=dropout)
@@ -72,7 +74,7 @@ def forward(self, inputs, memory_bank, src_pad_mask, tgt_pad_mask,
             query, attn = self.self_attn(input_norm, input_norm, input_norm,
                                          mask=dec_mask,
                                          layer_cache=layer_cache,
-                                         type="self")
+                                         attn_type="self")
         elif isinstance(self.self_attn, AverageAttention):
             query, attn = self.self_attn(input_norm, mask=dec_mask,
                                          layer_cache=layer_cache, step=step)
@@ -83,7 +85,7 @@ def forward(self, inputs, memory_bank, src_pad_mask, tgt_pad_mask,
         mid, attn = self.context_attn(memory_bank, memory_bank, query_norm,
                                       mask=src_pad_mask,
                                       layer_cache=layer_cache,
-                                      type="context")
+                                      attn_type="context")
         output = self.feed_forward(self.drop(mid) + query)
 
         return output, attn
@@ -127,7 +129,7 @@ class TransformerDecoder(DecoderBase):
 
     def __init__(self, num_layers, d_model, heads, d_ff,
                  copy_attn, self_attn_type, dropout, embeddings,
-                 max_relative_positions):
+                 max_relative_positions, aan_useffn):
         super(TransformerDecoder, self).__init__()
 
         self.embeddings = embeddings
@@ -138,7 +140,8 @@ def __init__(self, num_layers, d_model, heads, d_ff,
         self.transformer_layers = nn.ModuleList(
             [TransformerDecoderLayer(d_model, heads, d_ff, dropout,
              self_attn_type=self_attn_type,
-             max_relative_positions=max_relative_positions)
+             max_relative_positions=max_relative_positions,
+             aan_useffn=aan_useffn)
              for i in range(num_layers)])
 
         # previously, there was a GlobalAttention module here for copy
@@ -159,7 +162,8 @@ def from_opt(cls, opt, embeddings):
             opt.self_attn_type,
             opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
             embeddings,
-            opt.max_relative_positions)
+            opt.max_relative_positions,
+            opt.aan_useffn)
 
     def init_state(self, src, memory_bank, enc_hidden):
         """Initialize decoder state."""
@@ -233,7 +237,8 @@ def _init_cache(self, memory_bank):
         for i, layer in enumerate(self.transformer_layers):
             layer_cache = {"memory_keys": None, "memory_values": None}
             if isinstance(layer.self_attn, AverageAttention):
-                layer_cache["prev_g"] = torch.zeros((batch_size, 1, depth))
+                layer_cache["prev_g"] = torch.zeros((batch_size, 1, depth),
+                                                    device=memory_bank.device)
             else:
                 layer_cache["self_keys"] = None
                 layer_cache["self_values"] = None

onmt/encoders/transformer.py

@@ -46,7 +46,7 @@ def forward(self, inputs, mask):
         """
         input_norm = self.layer_norm(inputs)
         context, _ = self.self_attn(input_norm, input_norm, input_norm,
-                                    mask=mask, type="self")
+                                    mask=mask, attn_type="self")
         out = self.dropout(context) + inputs
         return self.feed_forward(out)
 

onmt/modules/average_attn.py

@@ -19,16 +19,16 @@ class AverageAttention(nn.Module):
        dropout (float): dropout parameter
     """
 
-    def __init__(self, model_dim, dropout=0.1):
+    def __init__(self, model_dim, dropout=0.1, aan_useffn=False):
         self.model_dim = model_dim
-
+        self.aan_useffn = aan_useffn
         super(AverageAttention, self).__init__()
-
-        self.average_layer = PositionwiseFeedForward(model_dim, model_dim,
-                                                     dropout)
+        if aan_useffn:
+            self.average_layer = PositionwiseFeedForward(model_dim, model_dim,
+                                                         dropout)
         self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)
 
-    def cumulative_average_mask(self, batch_size, inputs_len):
+    def cumulative_average_mask(self, batch_size, inputs_len, device):
         """
         Builds the mask to compute the cumulative average as described in
         :cite:`DBLP:journals/corr/abs-1805-00631` -- Figure 3
@@ -43,9 +43,10 @@ def cumulative_average_mask(self, batch_size, inputs_len):
             * A Tensor of shape ``(batch_size, input_len, input_len)``
         """
 
-        triangle = torch.tril(torch.ones(inputs_len, inputs_len))
-        weights = torch.ones(1, inputs_len) / torch.arange(
-            1, inputs_len + 1, dtype=torch.float)
+        triangle = torch.tril(torch.ones(inputs_len, inputs_len,
+                              dtype=torch.float, device=device))
+        weights = torch.ones(1, inputs_len, dtype=torch.float, device=device) \
+            / torch.arange(1, inputs_len + 1, dtype=torch.float, device=device)
         mask = triangle * weights.transpose(0, 1)
 
         return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)
@@ -72,14 +73,13 @@ def cumulative_average(self, inputs, mask_or_step,
 
         if layer_cache is not None:
             step = mask_or_step
-            device = inputs.device
             average_attention = (inputs + step *
-                                 layer_cache["prev_g"].to(device)) / (step + 1)
+                                 layer_cache["prev_g"]) / (step + 1)
             layer_cache["prev_g"] = average_attention
             return average_attention
         else:
             mask = mask_or_step
-            return torch.matmul(mask, inputs)
+            return torch.matmul(mask.to(inputs.dtype), inputs)
 
     def forward(self, inputs, mask=None, layer_cache=None, step=None):
         """
@@ -96,13 +96,12 @@ def forward(self, inputs, mask=None, layer_cache=None, step=None):
 
         batch_size = inputs.size(0)
         inputs_len = inputs.size(1)
-
-        device = inputs.device
         average_outputs = self.cumulative_average(
           inputs, self.cumulative_average_mask(batch_size,
-                                               inputs_len).to(device).float()
+                                               inputs_len, inputs.device)
           if layer_cache is None else step, layer_cache=layer_cache)
-        average_outputs = self.average_layer(average_outputs)
+        if self.aan_useffn:
+            average_outputs = self.average_layer(average_outputs)
         gating_outputs = self.gating_layer(torch.cat((inputs,
                                                       average_outputs), -1))
         input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)

onmt/modules/multi_headed_attn.py

@@ -75,7 +75,7 @@ def __init__(self, head_count, model_dim, dropout=0.1,
                 vocab_size, self.dim_per_head)
 
     def forward(self, key, value, query, mask=None,
-                layer_cache=None, type=None):
+                layer_cache=None, attn_type=None):
         """
         Compute the context vector and the attention vectors.
 
@@ -117,7 +117,6 @@ def forward(self, key, value, query, mask=None,
         head_count = self.head_count
         key_len = key.size(1)
         query_len = query.size(1)
-        device = key.device
 
         def shape(x):
             """Projection."""
@@ -131,23 +130,23 @@ def unshape(x):
 
         # 1) Project key, value, and query.
         if layer_cache is not None:
-            if type == "self":
+            if attn_type == "self":
                 query, key, value = self.linear_query(query),\
                                     self.linear_keys(query),\
                                     self.linear_values(query)
                 key = shape(key)
                 value = shape(value)
                 if layer_cache["self_keys"] is not None:
                     key = torch.cat(
-                        (layer_cache["self_keys"].to(device), key),
+                        (layer_cache["self_keys"], key),
                         dim=2)
                 if layer_cache["self_values"] is not None:
                     value = torch.cat(
-                        (layer_cache["self_values"].to(device), value),
+                        (layer_cache["self_values"], value),
                         dim=2)
                 layer_cache["self_keys"] = key
                 layer_cache["self_values"] = value
-            elif type == "context":
+            elif attn_type == "context":
                 query = self.linear_query(query)
                 if layer_cache["memory_keys"] is None:
                     key, value = self.linear_keys(key),\
@@ -166,18 +165,18 @@ def unshape(x):
             key = shape(key)
             value = shape(value)
 
-        if self.max_relative_positions > 0 and type == "self":
+        if self.max_relative_positions > 0 and attn_type == "self":
             key_len = key.size(2)
             # 1 or key_len x key_len
             relative_positions_matrix = generate_relative_positions_matrix(
                 key_len, self.max_relative_positions,
                 cache=True if layer_cache is not None else False)
             #  1 or key_len x key_len x dim_per_head
             relations_keys = self.relative_positions_embeddings(
-                relative_positions_matrix.to(device))
+                relative_positions_matrix.to(key.device))
             #  1 or key_len x key_len x dim_per_head
             relations_values = self.relative_positions_embeddings(
-                relative_positions_matrix.to(device))
+                relative_positions_matrix.to(key.device))
 
         query = shape(query)
 
@@ -189,7 +188,7 @@ def unshape(x):
         # batch x num_heads x query_len x key_len
         query_key = torch.matmul(query, key.transpose(2, 3))
 
-        if self.max_relative_positions > 0 and type == "self":
+        if self.max_relative_positions > 0 and attn_type == "self":
             scores = query_key + relative_matmul(query, relations_keys, True)
         else:
             scores = query_key
@@ -205,7 +204,7 @@ def unshape(x):
 
         context_original = torch.matmul(drop_attn, value)
 
-        if self.max_relative_positions > 0 and type == "self":
+        if self.max_relative_positions > 0 and attn_type == "self":
             context = unshape(context_original
                               + relative_matmul(drop_attn,
                                                 relations_values,

onmt/opts.py

@@ -151,6 +151,8 @@ def model_opts(parser):
               help='Number of heads for transformer self-attention')
     group.add('--transformer_ff', '-transformer_ff', type=int, default=2048,
               help='Size of hidden transformer feed-forward')
+    group.add('--aan_useffn', '-aan_useffn', action="store_true",
+              help='Turn on the FFN layer in the AAN decoder')
 
     # Generator and loss options.
     group.add('--copy_attn', '-copy_attn', action="store_true",

onmt/tests/pull_request_chk.sh

@@ -10,7 +10,7 @@ echo > ${LOG_FILE} # Empty the log file.
 PROJECT_ROOT=`dirname "$0"`"/../../"
 DATA_DIR="$PROJECT_ROOT/data"
 TEST_DIR="$PROJECT_ROOT/onmt/tests"
-PYTHON="python"
+PYTHON="python3"
 
 clean_up()
 {