Repository: OpenNMT/OpenNMT-tf
Pull request #327 (2019-02-09 10:24:35)

Add v2 Transformer encoder


Patch:
diff --git a/opennmt/layers/position.py b/opennmt/layers/position.py
index 726b060b..b3dbb74f 100644
--- a/opennmt/layers/position.py
+++ b/opennmt/layers/position.py
@@ -67,14 +67,13 @@ class PositionEncoder(object):
     timesteps = tf.shape(inputs)[1]
     input_dim = inputs.get_shape().as_list()[-1]
 
-    with tf.variable_scope("position_encoding"):
-      if position is None:
-        positions = tf.range(timesteps) + 1
-      else:
-        positions = [position]
-      position_encoding = self.encode([positions], input_dim, dtype=inputs.dtype)
-      position_encoding = tf.tile(position_encoding, [batch_size, 1, 1])
-      return self.reducer([inputs, position_encoding])
+    if position is None:
+      positions = tf.range(timesteps) + 1
+    else:
+      positions = [position]
+    position_encoding = self.encode([positions], input_dim, dtype=inputs.dtype)
+    position_encoding = tf.tile(position_encoding, [batch_size, 1, 1])
+    return self.reducer([inputs, position_encoding])
 
   def apply(self, inputs, sequence_length=None):
     """Shortcut for ``__call__``."""
@@ -136,8 +135,9 @@ class PositionEmbedder(PositionEncoder):
 
   def encode(self, positions, depth, dtype=tf.float32):
     positions = tf.minimum(positions, self.maximum_position)
-    embeddings = tf.get_variable(
-        "w_embs", shape=[self.maximum_position + 1, depth], dtype=dtype)
+    with tf.variable_scope("position_encoding"):
+      embeddings = tf.get_variable(
+          "w_embs", shape=[self.maximum_position + 1, depth], dtype=dtype)
     return tf.nn.embedding_lookup(embeddings, positions)
 
 
diff --git a/opennmt/utils/misc.py b/opennmt/utils/misc.py
index 77582459..d58fb081 100644
--- a/opennmt/utils/misc.py
+++ b/opennmt/utils/misc.py
@@ -81,6 +81,12 @@ def classes_in_module(module, public_only=False):
           if (inspect.isclass(getattr(module, symbol))
               and (not public_only or not symbol.startswith("_"))))
 
+def function_args(fun):
+  """Returns the name of :obj:`fun` arguments."""
+  if hasattr(inspect, "getfullargspec"):
+    return inspect.getfullargspec(fun).args
+  return inspect.getargspec(fun).args  # pylint: disable=deprecated-method
+
 def get_third_party_dir():
   """Returns a path to the third_party directory."""
   utils_dir = os.path.dirname(__file__)
diff --git a/opennmt/v2/__init__.py b/opennmt/v2/__init__.py
index 15a4fbb0..0bde83b6 100644
--- a/opennmt/v2/__init__.py
+++ b/opennmt/v2/__init__.py
@@ -1,3 +1,4 @@
 """Module for TensorFlow 2.0 compatibility."""
 
+from opennmt.v2 import encoders
 from opennmt.v2 import layers
diff --git a/opennmt/v2/encoders/__init__.py b/opennmt/v2/encoders/__init__.py
new file mode 100644
index 00000000..1ae3939b
--- /dev/null
+++ b/opennmt/v2/encoders/__init__.py
@@ -0,0 +1,4 @@
+"""Encoders."""
+
+from opennmt.v2.encoders.encoder import Encoder
+from opennmt.v2.encoders.self_attention_encoder import SelfAttentionEncoder
diff --git a/opennmt/v2/encoders/encoder.py b/opennmt/v2/encoders/encoder.py
new file mode 100644
index 00000000..014c8203
--- /dev/null
+++ b/opennmt/v2/encoders/encoder.py
@@ -0,0 +1,12 @@
+"""Base class for encoders and generic multi encoders."""
+
+import abc
+import six
+
+import tensorflow as tf
+
+
+@six.add_metaclass(abc.ABCMeta)
+class Encoder(tf.keras.layers.Layer):
+  """Base class for encoders."""
+  pass
diff --git a/opennmt/v2/encoders/self_attention_encoder.py b/opennmt/v2/encoders/self_attention_encoder.py
new file mode 100644
index 00000000..77afe1e9
--- /dev/null
+++ b/opennmt/v2/encoders/self_attention_encoder.py
@@ -0,0 +1,107 @@
+"""Define the self-attention encoder."""
+
+import tensorflow as tf
+
+from opennmt.v2.encoders.encoder import Encoder
+from opennmt.v2.layers import common, transformer
+from opennmt.layers.position import SinusoidalPositionEncoder
+
+
+class SelfAttentionEncoder(Encoder):
+  """Encoder using self-attention as described in
+  https://arxiv.org/abs/1706.03762.
+  """
+
+  def __init__(self,
+               num_layers,
+               num_units=512,
+               num_heads=8,
+               ffn_inner_dim=2048,
+               dropout=0.1,
+               attention_dropout=0.1,
+               relu_dropout=0.1,
+               position_encoder=None,
+               **kwargs):
+    """Initializes the parameters of the encoder.
+
+    Args:
+      num_layers: The number of layers.
+      num_units: The number of hidden units.
+      num_heads: The number of heads in the multi-head attention.
+      ffn_inner_dim: The number of units of the inner linear transformation
+        in the feed forward layer.
+      dropout: The probability to drop units from the outputs.
+      attention_dropout: The probability to drop units from the attention.
+      relu_dropout: The probability to drop units from the ReLU activation in
+        the feed forward layer.
+      position_encoder: The :class:`opennmt.layers.position.PositionEncoder` to
+        apply on inputs. If ``None``, defaults to
+        :class:`opennmt.layers.position.SinusoidalPositionEncoder`.
+    """
+    super(SelfAttentionEncoder, self).__init__(**kwargs)
+    self.num_units = num_units
+    self.dropout = dropout
+    self.position_encoder = position_encoder
+    if self.position_encoder is None:
+      self.position_encoder = SinusoidalPositionEncoder()
+    self.layer_norm = common.LayerNorm()
+    self.layers = [
+        _SelfAttentionEncoderLayer(
+            num_units,
+            num_heads,
+            ffn_inner_dim,
+            dropout=dropout,
+            attention_dropout=attention_dropout,
+            relu_dropout=relu_dropout,
+            name="layer_%d" % i)
+        for i in range(num_layers)]
+
+  def call(self, inputs, mask=None, training=None):  # pylint: disable=arguments-differ
+    """Encodes :obj:`inputs`."""
+    inputs *= self.num_units**0.5
+    inputs = self.position_encoder(inputs)
+    if training:
+      inputs = tf.nn.dropout(inputs, rate=self.dropout)
+    for layer in self.layers:
+      inputs = layer(inputs, mask=mask, training=training)
+    outputs = self.layer_norm(inputs)
+    return outputs, None, mask
+
+
+class _SelfAttentionEncoderLayer(tf.keras.layers.Layer):
+  """Implements one self-attention encoding layer."""
+
+  def __init__(self,
+               num_units,
+               num_heads,
+               ffn_inner_dim,
+               dropout=0.1,
+               attention_dropout=0.1,
+               relu_dropout=0.1,
+               **kwargs):
+    """Initializes the layer.
+
+    Args:
+      num_units: The number of hidden units.
+      num_heads: The number of heads in the multi-head attention.
+      ffn_inner_dim: The number of units of the inner linear transformation
+        in the feed forward layer.
+      dropout: The probability to drop units from the outputs.
+      attention_dropout: The probability to drop units from the attention.
+      relu_dropout: The probability to drop units from the ReLU activation in
+        the feed forward layer.
+      kwargs: Additional layer arguments.
+    """
+    super(_SelfAttentionEncoderLayer, self).__init__(**kwargs)
+    self.self_attention = transformer.MultiHeadAttention(
+        num_heads, num_units, dropout=attention_dropout)
+    self.self_attention = common.LayerWrapper(self.self_attention, dropout=dropout)
+    self.ffn = transformer.FeedForwardNetwork(
+        ffn_inner_dim, num_units, dropout=relu_dropout)
+    self.ffn = common.LayerWrapper(self.ffn, dropout=dropout)
+
+  def call(self, x, mask=None, training=None):  # pylint: disable=arguments-differ
+    """Runs the encoder layer."""
+    y, _ = self.self_attention(x, mask=mask, training=training)
+    y = self.ffn(y, training=training)
+    return y
diff --git a/opennmt/v2/layers/__init__.py b/opennmt/v2/layers/__init__.py
index 2380eea0..4c81ef2e 100644
--- a/opennmt/v2/layers/__init__.py
+++ b/opennmt/v2/layers/__init__.py
@@ -1,4 +1,4 @@
 """Layers module."""
 
-from opennmt.v2.layers.common import LayerNorm
+from opennmt.v2.layers.common import LayerNorm, LayerWrapper
 from opennmt.v2.layers.transformer import FeedForwardNetwork, MultiHeadAttention
diff --git a/opennmt/v2/layers/common.py b/opennmt/v2/layers/common.py
index b5efc9f7..9dcdcfe9 100644
--- a/opennmt/v2/layers/common.py
+++ b/opennmt/v2/layers/common.py
@@ -2,6 +2,8 @@
 
 import tensorflow as tf
 
+from opennmt.utils.misc import function_args
+
 
 class LayerNorm(tf.keras.layers.Layer):
   """Layer normalization."""
@@ -31,3 +33,55 @@ class LayerNorm(tf.keras.layers.Layer):
     variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
     norm_x = (x - mean) * tf.math.rsqrt(variance + self.epsilon)
     return norm_x * self.scale + self.bias
+
+
+class LayerWrapper(tf.keras.layers.Wrapper):
+  """Layer wrapper for input/output normalization and residual connections."""
+
+  def __init__(self,
+               layer,
+               normalize_input=True,
+               normalize_output=False,
+               residual_connections=True,
+               dropout=0.1,
+               **kwargs):
+    """Initializes the layer.
+
+    Args:
+      layer: The layer to wrap.
+      normalize_input: Apply layer normalization on the input.
+      normalize_output: Apply layer normalization on the output.
+      residual_connections: Add the inputs to layer outputs (if their shape are
+        compatible).
+      dropout: The probability to drop units in the layer output.
+      kwargs: Additional layer arguments.
+    """
+    super(LayerWrapper, self).__init__(layer, **kwargs)
+    self.input_layer_norm = LayerNorm() if normalize_input else None
+    self.output_layer_norm = LayerNorm() if normalize_output else None
+    self.residual_connections = residual_connections
+    self.dropout = dropout
+
+  def call(self, inputs, training=None, **kwargs):  # pylint: disable=arguments-differ
+    """Runs the wrapper."""
+    x = inputs
+    if self.input_layer_norm is not None:
+      x = self.input_layer_norm(x)
+    if "training" in function_args(self.layer.call):
+      kwargs["training"] = training
+    all_outputs = self.layer(x, **kwargs)
+    if isinstance(all_outputs, tuple):
+      outputs = all_outputs[0]
+      extra_outputs = list(all_outputs)[1:]
+    else:
+      outputs = all_outputs
+      extra_outputs = None
+    if training:
+      outputs = tf.nn.dropout(outputs, rate=self.dropout)
+    if self.residual_connections and outputs.shape[-1] == inputs.shape[-1]:
+      outputs += inputs
+    if self.output_layer_norm is not None:
+      outputs = self.output_layer_norm(outputs)
+    if extra_outputs:
+      return tuple([outputs] + extra_outputs)
+    return outputs
diff --git a/opennmt/v2/layers/transformer.py b/opennmt/v2/layers/transformer.py
index 6325398c..4e525ee0 100644
--- a/opennmt/v2/layers/transformer.py
+++ b/opennmt/v2/layers/transformer.py
@@ -118,7 +118,7 @@ class MultiHeadAttention(tf.keras.layers.Layer):
       mask = tf.expand_dims(tf.cast(mask, tf.float32), 1)  # Broadcast on heads dimension.
       dot = tf.cast(tf.cast(dot, tf.float32) * mask + ((1.0 - mask) * tf.float32.min), dot.dtype)
     attn = tf.cast(tf.nn.softmax(tf.cast(dot, tf.float32)), dot.dtype)
-    drop_attn = tf.nn.dropout(attn, dropout) if training else attn
+    drop_attn = tf.nn.dropout(attn, self.dropout) if training else attn
     heads = tf.matmul(drop_attn, values)
 
     # Concatenate all heads output.


Test patch:
diff --git a/opennmt/v2/tests/common_layers_test.py b/opennmt/v2/tests/common_layers_test.py
index 6236ee6d..898fbe5c 100644
--- a/opennmt/v2/tests/common_layers_test.py
+++ b/opennmt/v2/tests/common_layers_test.py
@@ -11,6 +11,18 @@ class CommonLayersTest(tf.test.TestCase):
     y = layer_norm(x)
     self.assertEqual(y.shape, x.shape)
 
+  def testLayerWrapper(self):
+    layer = common.LayerWrapper(tf.keras.layers.Dense(10))
+    x = tf.random.uniform([4, 5, 10])
+    y = layer(x)
+    self.assertEqual(y.shape, x.shape)
+
+  def testLayerWrapperInputOutputDepthMismatch(self):
+    layer = common.LayerWrapper(tf.keras.layers.Dense(10))
+    x = tf.random.uniform([4, 5, 5])
+    y = layer(x)
+    self.assertListEqual(y.shape.as_list(), [4, 5, 10])
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/opennmt/v2/tests/encoder_test.py b/opennmt/v2/tests/encoder_test.py
new file mode 100644
index 00000000..eac84a73
--- /dev/null
+++ b/opennmt/v2/tests/encoder_test.py
@@ -0,0 +1,22 @@
+from parameterized import parameterized
+
+import tensorflow as tf
+
+from opennmt.v2 import encoders
+
+
+class EncoderTest(tf.test.TestCase):
+
+  @parameterized.expand([[tf.float32], [tf.float16]])
+  def testSelfAttentionEncoder(self, dtype):
+    encoder = encoders.SelfAttentionEncoder(3, num_units=20, num_heads=4, ffn_inner_dim=40)
+    x = tf.random.uniform([4, 5, 10], dtype=dtype)
+    mask = tf.expand_dims(tf.sequence_mask([4, 3, 5, 2]), 1)
+    y, _, output_mask = encoder(x, mask=mask, training=True)
+    self.assertListEqual(y.shape.as_list(), [4, 5, 20])
+    self.assertEqual(y.dtype, dtype)
+    self.assertAllEqual(output_mask, mask)
+
+
+if __name__ == "__main__":
+  tf.test.main()
diff --git a/opennmt/v2/tests/transformer_test.py b/opennmt/v2/tests/transformer_test.py
index 3b24d636..259d1958 100644
--- a/opennmt/v2/tests/transformer_test.py
+++ b/opennmt/v2/tests/transformer_test.py
@@ -63,6 +63,17 @@ class TransformerTest(tf.test.TestCase):
     y2, cache = attention(x, memory=memory, mask=mask, cache=cache)
     self.assertAllEqual(y1, y2)
 
+  def testMultiHeadAttentionMask(self):
+    attention = transformer.MultiHeadAttention(4, 20, return_attention=True)
+    queries = tf.random.uniform([4, 5, 10])
+    memory = tf.random.uniform([4, 3, 10])
+    mask = tf.expand_dims(tf.sequence_mask([1, 3, 2, 2]), 1)
+    _, _, attention = attention(queries, memory=memory, mask=mask)
+    attention = tf.reshape(attention, [4, -1, 3])
+    mask = tf.broadcast_to(mask, attention.shape)
+    padding = tf.boolean_mask(attention, tf.logical_not(mask))
+    self.assertAllEqual(tf.reduce_sum(padding), 0)
+
 
 if __name__ == "__main__":
   tf.test.main()