Added example model

Author: EndingCredits

tensor2tensor/models/transformer_alternative_.py renamed to tensor2tensor/models/transformer_alternative.py

@@ -60,6 +60,12 @@ def model_fn_body(self, features):
         transformer_prepare_encoder(inputs, target_space, hparams) )
     (decoder_input, decoder_self_attention_bias) = transformer.\
         transformer_prepare_decoder(targets, hparams)
+    
+    # We need masks of the form batch size x input sequences
+    # Biases seem to be of the form batch_size x 1 x input sequences x vec dim
+    #  Squeeze out dim one, and get the first element of each vector
+    encoder_mask = tf.squeeze(encoder_attention_bias, [1])[:,:,0]
+    decoder_mask = tf.squeeze(decoder_self_attention_bias, [1])[:,:,0]
 
     def residual_fn(x, y):
       return common_layers.layer_norm(x + tf.nn.dropout(
@@ -68,35 +74,60 @@ def residual_fn(x, y):
     encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.residual_dropout)
     decoder_input = tf.nn.dropout(decoder_input, 1.0 - hparams.residual_dropout)
     encoder_output = alt_transformer_encoder(
-        encoder_input, residual_fn, encoder_attention_bias, hparams)
+        encoder_input, residual_fn, encoder_mask, hparams)
 
     decoder_output = alt_transformer_decoder(
-        decoder_input, encoder_output, residual_fn, decoder_self_attention_bias,
+        decoder_input, encoder_output, residual_fn, decoder_mask,
         encoder_attention_bias, hparams)
 
     decoder_output = tf.expand_dims(decoder_output, 2)
 
     return decoder_output
 
 
+    
+def composite_layer(inputs, mask, hparams):
+  x = inputs
+  
+  # Applies ravanbakhsh on top of each other
+  if hparams.composite_layer_type == "ravanbakhsh":
+    for layer in xrange(hparams.layers_per_layer):
+      with tf.variable_scope(".%d" % layer):
+        x = common_layers.ravanbakhsh_set_layer(
+                hparams.hidden_size,
+                x,
+                mask=mask,
+                dropout=0.0)
+  
+  # Transforms elements to get a context, and then uses this in a final layer
+  elif hparams.composite_layer_type == "reembedding":
+    initial_elems = x
+    # Transform elements n times and then pool
+    for layer in xrange(hparams.layers_per_layer):
+      with tf.variable_scope(".%d" % layer):
+        x = common_layers.linear_set_layer(
+                hparams.hidden_size,
+                x,
+                dropout=0.0)
+    context = common_layers.global_pool_1d(x, mask=mask)
+     
+    #Final layer
+    x = common_layers.linear_set_layer(
+            hparams.hidden_size,
+            x,
+            context=context,
+            dropout=0.0)
+                 
+  return x
+    
+
+
 def alt_transformer_encoder(encoder_input,
                             residual_fn,
-                            encoder_attention_bias,
+                            mask,
                             hparams,
                             name="encoder"):
-  """
-  A stack of transformer layers.
 
-  Args:
-    encoder_input: a Tensor
-    residual_fn: a function from (layer_input, layer_output) -> combined_output
-
-    hparams: hyperparameters for model
-    name: a string
-
-  Returns:
-    y: a Tensors
-  """
   x = encoder_input
 
   # Summaries don't work in multi-problem setting yet.
@@ -105,36 +136,19 @@ def alt_transformer_encoder(encoder_input,
   with tf.variable_scope(name):
     for layer in xrange(hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
-        x = residual_fn(
-            x,
-            ravanbakhsh_set_layer(hparams.hidden_size, x, mask=encoder_attention_bias)
-        )
+        x = residual_fn(x, composite_layer(x, mask, hparams))
 
   return x
 
 
 def alt_transformer_decoder(decoder_input,
                             encoder_output,
                             residual_fn,
-                            decoder_self_attention_bias,
+                            mask,
                             encoder_decoder_attention_bias,
                             hparams,
                             name="decoder"):
-  """
-  A stack of transformer layers.
-
-  Args:
-    decoder_input: a Tensor
-    encoder_output: a Tensor
-    residual_fn: a function from (layer_input, layer_output) -> combined_output
-    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
-      (see common_attention.attention_bias())
-    hparams: hyperparameters for model
-    name: a string
-
-  Returns:
-    y: a Tensors
-  """
+
   x = decoder_input
 
   # Summaries don't work in multi-problem setting yet.
@@ -143,30 +157,34 @@ def alt_transformer_decoder(decoder_input,
     for layer in xrange(hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
 
-        x = residual_fn(
-            x,
-            ravanbakhsh_set_layer(hparams.hidden_size,
-                common_attention.multihead_attention(
-                    x,
-                    encoder_output,
-                    encoder_decoder_attention_bias,
-                    hparams.attention_key_channels or hparams.hidden_size,
-                    hparams.attention_value_channels or hparams.hidden_size,
-                    hparams.hidden_size,
-                    hparams.num_heads,
-                    hparams.attention_dropout,
-                    summaries=summaries,
-                    name="encdec_attention"),
-            mask=decoder_self_attention_bias)
-        )
+        x_ = common_attention.multihead_attention(
+                 x,
+                 encoder_output,
+                 encoder_decoder_attention_bias,
+                 hparams.attention_key_channels or hparams.hidden_size,
+                 hparams.attention_value_channels or hparams.hidden_size,
+                 hparams.hidden_size,
+                 hparams.num_heads,
+                 hparams.attention_dropout,
+                 summaries=summaries,
+                 name="encdec_attention")
+
+        x_ = residual_fn(x_, composite_layer(x_, mask, hparams))
+        x = residual_fn(x, x_)
 
   return x
 
 
+
+
+
 @registry.register_hparams
 def transformer_alt():
   """Set of hyperparameters."""
   hparams = transformer.transformer_base()
+  hparams.batch_size = 64
   hparams.add_hparam("layers_per_layer", 4)
+  #hparams.add_hparam("composite_layer_type", "ravanbakhsh") #ravanbakhsh or reembedding
+  hparams.add_hparam("composite_layer_type", "reembedding")
   return hparams