Clean up autoencoder code, correct modality bug, add autoregressive baselines and tests.

PiperOrigin-RevId: 193435665

Author: Lukasz Kaiser

tensor2tensor/layers/common_layers.py

@@ -72,6 +72,8 @@ def dropout_with_broadcast_dims(x, keep_prob, broadcast_dims=None, **kwargs):
   if broadcast_dims:
     shape = tf.shape(x)
     ndims = len(x.get_shape())
+    # Allow dimensions like "-1" as well.
+    broadcast_dims = [dim + ndims if dim < 0 else dim for dim in broadcast_dims]
     kwargs["noise_shape"] = [
         1 if i in broadcast_dims else shape[i] for i in xrange(ndims)]
   return tf.nn.dropout(x, keep_prob, **kwargs)
@@ -441,7 +443,7 @@ def conv2d_kernel(kernel_size_arg, name_suffix):
   return conv2d_kernel(kernel_size, "single")
 
 
-def conv(inputs, filters, kernel_size, dilation_rate=1, **kwargs):
+def conv(inputs, filters, kernel_size, dilation_rate=(1, 1), **kwargs):
   return conv_internal(
       tf.layers.conv2d,
       inputs,

tensor2tensor/layers/modalities.py

@@ -258,8 +258,7 @@ def loss(self, logits, targets):
         logits,
         targets,
         self._model_hparams.label_smoothing,
-        weights_fn=self.targets_weights_fn,
-        gaussian=True)
+        weights_fn=self.targets_weights_fn)
 
 
 @registry.register_image_modality("image_channel_compress")
@@ -535,8 +534,7 @@ def loss(self, logits, targets):
         logits,
         targets,
         self._model_hparams.label_smoothing,
-        weights_fn=self.targets_weights_fn,
-        gaussian=True)
+        weights_fn=self.targets_weights_fn)
 
 
 @registry.register_class_label_modality("default")

tensor2tensor/models/basic.py

@@ -206,7 +206,7 @@ def basic_autoencoder():
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup"
-  hparams.label_smoothing = 0.05
+  hparams.label_smoothing = 0.0
   hparams.batch_size = 128
   hparams.hidden_size = 64
   hparams.num_hidden_layers = 5

tensor2tensor/models/basic_test.py

@@ -0,0 +1,70 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Basic nets tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+import numpy as np
+
+from tensor2tensor.data_generators import all_problems  # pylint: disable=unused-import
+from tensor2tensor.models import basic
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+
+class BasicTest(tf.test.TestCase):
+
+  def testBasicFcRelu(self):
+    x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
+    y = np.random.random_integers(0, high=9, size=(1, 1))
+    hparams = trainer_lib.create_hparams(
+        "basic_fc_small", problem_name="image_mnist", data_dir=".")
+    with self.test_session() as session:
+      features = {
+          "inputs": tf.constant(x, dtype=tf.int32),
+          "targets": tf.constant(y, dtype=tf.int32),
+      }
+      model = basic.BasicFcRelu(hparams, tf.estimator.ModeKeys.TRAIN)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    self.assertEqual(res.shape, (1, 1, 1, 1, 10))
+
+  def testBasicAutoencoder(self):
+    x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
+    y = np.random.random_integers(0, high=9, size=(1, 1))
+    hparams = trainer_lib.create_hparams(
+        "basic_autoencoder", problem_name="image_mnist_rev", data_dir=".")
+    with self.test_session() as session:
+      features = {
+          "targets": tf.constant(x, dtype=tf.int32),
+          "inputs": tf.constant(y, dtype=tf.int32),
+      }
+      tf.train.create_global_step()
+      model = basic.BasicAutoencoder(hparams, tf.estimator.ModeKeys.TRAIN)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    self.assertEqual(res.shape, (1, 28, 28, 1, 256))
+
+
+if __name__ == "__main__":
+  tf.test.main()

tensor2tensor/models/research/autoencoders.py

@@ -30,7 +30,53 @@
 
 
 @registry.register_model
-class ResidualAutoencoder(basic.BasicAutoencoder):
+class AutoencoderAutoregressive(basic.BasicAutoencoder):
+  """Autoencoder with an autoregressive part."""
+
+  def body(self, features):
+    hparams = self._hparams
+    shape = common_layers.shape_list(features["targets"])
+    # Run the basic autoencoder part first.
+    basic_result, losses = super(AutoencoderAutoregressive, self).body(features)
+    # Prepare inputs for autoregressive modes.
+    targets_keep_prob = 1.0 - hparams.autoregressive_dropout
+    targets_dropout = common_layers.dropout_with_broadcast_dims(
+        features["targets"], targets_keep_prob, broadcast_dims=[-1])
+    targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[3]])
+    targets_shifted = common_layers.shift_right_3d(targets1d)
+    basic1d = tf.reshape(basic_result, [shape[0], -1, shape[3]])
+    concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
+    # The forget_base hparam sets purely-autoregressive mode, no autoencoder.
+    if hparams.autoregressive_forget_base:
+      concat1d = tf.reshape(features["targets"], [shape[0], -1, shape[3]])
+      concat1d = common_layers.shift_right_3d(concat1d)
+    # The autoregressive part depends on the mode.
+    if hparams.autoregressive_mode == "none":
+      assert not hparams.autoregressive_forget_base
+      return basic_result, losses
+    if hparams.autoregressive_mode == "conv3":
+      res = common_layers.conv1d(concat1d, shape[3], 3, padding="LEFT",
+                                 activation=common_layers.belu,
+                                 name="autoregressive_conv3")
+      return tf.reshape(res, shape), losses
+    if hparams.autoregressive_mode == "conv5":
+      res = common_layers.conv1d(concat1d, shape[3], 5, padding="LEFT",
+                                 activation=common_layers.belu,
+                                 name="autoregressive_conv5")
+      return tf.reshape(res, shape), losses
+    if hparams.autoregressive_mode == "sru":
+      res = common_layers.conv1d(concat1d, shape[3], 3, padding="LEFT",
+                                 activation=common_layers.belu,
+                                 name="autoregressive_sru_conv3")
+      res = common_layers.sru(res)
+      return tf.reshape(res, shape), losses
+
+    raise ValueError("Unsupported autoregressive mode: %s"
+                     % hparams.autoregressive_mode)
+
+
+@registry.register_model
+class AutoencoderResidual(AutoencoderAutoregressive):
   """Residual autoencoder."""
 
   def encoder(self, x):
@@ -106,7 +152,7 @@ def decoder(self, x):
 
 
 @registry.register_model
-class BasicDiscreteAutoencoder(basic.BasicAutoencoder):
+class AutoencoderBasicDiscrete(AutoencoderAutoregressive):
   """Discrete autoencoder."""
 
   def bottleneck(self, x):
@@ -132,7 +178,7 @@ def sample(self):
 
 
 @registry.register_model
-class ResidualDiscreteAutoencoder(ResidualAutoencoder):
+class AutoencoderResidualDiscrete(AutoencoderResidual):
   """Discrete residual autoencoder."""
 
   def bottleneck(self, x, bottleneck_size=None):
@@ -160,13 +206,15 @@ def sample(self):
     size = [hp.batch_size, hp.sample_height // div_x, hp.sample_width // div_y,
             hp.bottleneck_size]
     rand = tf.random_uniform(size)
-    res1 = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
-    res2 = tf.zeros_like(rand) - 1.0
-    return tf.concat([res2[:, :, :, :2], res1[:, :, :, 2:]], axis=-1)
+    res = 2.0 * tf.to_float(tf.less(0.5, rand)) - 1.0
+    # If you want to set some first bits to a fixed value, do this:
+    # fixed = tf.zeros_like(rand) - 1.0
+    # res = tf.concat([fixed[:, :, :, :2], res[:, :, :, 2:]], axis=-1)
+    return res
 
 
 @registry.register_model
-class OrderedDiscreteAutoencoder(ResidualDiscreteAutoencoder):
+class AutoencoderOrderedDiscrete(AutoencoderResidualDiscrete):
   """Ordered discrete autoencoder."""
 
   def bottleneck(self, x):
@@ -195,7 +243,7 @@ def bottleneck(self, x):
 
 
 @registry.register_model
-class StackedAutoencoder(ResidualDiscreteAutoencoder):
+class AutoencoderStacked(AutoencoderResidualDiscrete):
   """A stacked autoencoder."""
 
   def stack(self, b, size, bottleneck_size, name):
@@ -290,9 +338,19 @@ def body(self, features):
 
 
 @registry.register_hparams
-def residual_autoencoder():
-  """Residual autoencoder model."""
+def autoencoder_autoregressive():
+  """Autoregressive autoencoder model."""
   hparams = basic.basic_autoencoder()
+  hparams.add_hparam("autoregressive_forget_base", False)
+  hparams.add_hparam("autoregressive_mode", "conv3")
+  hparams.add_hparam("autoregressive_dropout", 0.4)
+  return hparams
+
+
+@registry.register_hparams
+def autoencoder_residual():
+  """Residual autoencoder model."""
+  hparams = autoencoder_autoregressive()
   hparams.optimizer = "Adam"
   hparams.learning_rate_constant = 0.0001
   hparams.learning_rate_warmup_steps = 500
@@ -311,9 +369,9 @@ def residual_autoencoder():
 
 
 @registry.register_hparams
-def basic_discrete_autoencoder():
+def autoencoder_basic_discrete():
   """Basic autoencoder model."""
-  hparams = basic.basic_autoencoder()
+  hparams = autoencoder_autoregressive()
   hparams.num_hidden_layers = 5
   hparams.hidden_size = 64
   hparams.bottleneck_size = 4096
@@ -324,9 +382,9 @@ def basic_discrete_autoencoder():
 
 
 @registry.register_hparams
-def residual_discrete_autoencoder():
+def autoencoder_residual_discrete():
   """Residual discrete autoencoder model."""
-  hparams = residual_autoencoder()
+  hparams = autoencoder_residual()
   hparams.bottleneck_size = 4096
   hparams.bottleneck_noise = 0.1
   hparams.bottleneck_warmup_steps = 3000
@@ -339,9 +397,9 @@ def residual_discrete_autoencoder():
 
 
 @registry.register_hparams
-def residual_discrete_autoencoder_big():
+def autoencoder_residual_discrete_big():
   """Residual discrete autoencoder model, big version."""
-  hparams = residual_discrete_autoencoder()
+  hparams = autoencoder_residual_discrete()
   hparams.hidden_size = 128
   hparams.max_hidden_size = 4096
   hparams.bottleneck_noise = 0.1
@@ -351,15 +409,15 @@ def residual_discrete_autoencoder_big():
 
 
 @registry.register_hparams
-def ordered_discrete_autoencoder():
+def autoencoder_ordered_discrete():
   """Basic autoencoder model."""
-  hparams = residual_discrete_autoencoder()
+  hparams = autoencoder_residual_discrete()
   return hparams
 
 
 @registry.register_hparams
-def stacked_autoencoder():
+def autoencoder_stacked():
   """Stacked autoencoder model."""
-  hparams = residual_discrete_autoencoder()
+  hparams = autoencoder_residual_discrete()
   hparams.bottleneck_size = 128
   return hparams

tensor2tensor/models/research/autoencoders_test.py

@@ -0,0 +1,84 @@
+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Autoencoders tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+import numpy as np
+
+from tensor2tensor.data_generators import all_problems  # pylint: disable=unused-import
+from tensor2tensor.models.research import autoencoders  # pylint: disable=unused-import
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import trainer_lib
+
+import tensorflow as tf
+
+
+class AutoencoderTest(tf.test.TestCase):
+
+  def getMnistRandomOutput(self, model_name, hparams_set=None,
+                           mode=tf.estimator.ModeKeys.TRAIN):
+    hparams_set = hparams_set or model_name
+    x = np.random.random_integers(0, high=255, size=(1, 28, 28, 1))
+    y = np.random.random_integers(0, high=9, size=(1, 1))
+    hparams = trainer_lib.create_hparams(
+        hparams_set, problem_name="image_mnist_rev", data_dir=".")
+    with self.test_session() as session:
+      features = {
+          "targets": tf.constant(x, dtype=tf.int32),
+          "inputs": tf.constant(y, dtype=tf.int32),
+      }
+      tf.train.create_global_step()
+      model = registry.model(model_name)(hparams, mode)
+      logits, _ = model(features)
+      session.run(tf.global_variables_initializer())
+      res = session.run(logits)
+    return res
+
+  @property
+  def mnistOutputShape(self):
+    return (1, 28, 28, 1, 256)
+
+  def testAutoencoderAutoregressive(self):
+    res = self.getMnistRandomOutput("autoencoder_autoregressive")
+    self.assertEqual(res.shape, self.mnistOutputShape)
+
+  def testAutoencoderResidual(self):
+    res = self.getMnistRandomOutput("autoencoder_residual")
+    self.assertEqual(res.shape, self.mnistOutputShape)
+
+  def testAutoencoderBasicDiscrete(self):
+    res = self.getMnistRandomOutput("autoencoder_basic_discrete")
+    self.assertEqual(res.shape, self.mnistOutputShape)
+
+  def testAutoencoderResidualDiscrete(self):
+    res = self.getMnistRandomOutput("autoencoder_residual_discrete")
+    self.assertEqual(res.shape, self.mnistOutputShape)
+
+  def testAutoencoderOrderedDiscrete(self):
+    res = self.getMnistRandomOutput("autoencoder_ordered_discrete")
+    self.assertEqual(res.shape, self.mnistOutputShape)
+
+  def testAutoencoderStacked(self):
+    res = self.getMnistRandomOutput("autoencoder_stacked")
+    self.assertEqual(res.shape, self.mnistOutputShape)
+
+if __name__ == "__main__":
+  tf.test.main()