Add sampling with temperature and cifar10 8 by 8 dataset.

PiperOrigin-RevId: 172031867

Author: Ashish Vaswani

tensor2tensor/data_generators/image.py

@@ -42,6 +42,12 @@
 import tensorflow as tf
 
 
+def resize_by_area(img, size):
+  """image resize function used by quite a few image problems."""
+  return tf.to_int64(
+      tf.image.resize_images(img, [size, size], tf.image.ResizeMethod.AREA))
+
+
 class ImageProblem(problem.Problem):
 
   def example_reading_spec(self, label_key=None):
@@ -93,16 +99,12 @@ class ImageCeleba(ImageProblem):
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
 
-    def resize(img, size):
-      return tf.to_int64(
-          tf.image.resize_images(img, [size, size], tf.image.ResizeMethod.AREA))
-
     inputs = example["inputs"]
     # Remove boundaries in CelebA images. Remove 40 pixels each side
     # vertically and 20 pixels each side horizontally.
     inputs = tf.image.crop_to_bounding_box(inputs, 40, 20, 218 - 80, 178 - 40)
-    example["inputs"] = resize(inputs, 8)
-    example["targets"] = resize(inputs, 32)
+    example["inputs"] = resize_by_area(inputs, 8)
+    example["targets"] = resize_by_area(inputs, 32)
     return example
 
   def hparams(self, defaults, unused_model_hparams):
@@ -388,14 +390,10 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
 
-    def resize(img, size):
-      return tf.to_int64(
-          tf.image.resize_images(img, [size, size], tf.image.ResizeMethod.AREA))
-
     inputs = example["inputs"]
     # For Img2Img resize input and output images as desired.
-    example["inputs"] = resize(inputs, 8)
-    example["targets"] = resize(inputs, 32)
+    example["inputs"] = resize_by_area(inputs, 8)
+    example["targets"] = resize_by_area(inputs, 32)
     return example
 
   def hparams(self, defaults, unused_model_hparams):
@@ -654,6 +652,18 @@ def preprocess_example(self, example, mode, unused_hparams):
     return example
 
 
+@registry.register_problem
+class ImageCifar10Plain8(ImageCifar10):
+  """CIFAR-10 rescaled to 8x8 for output: Conditional image generation."""
+
+  def dataset_filename(self):
+    return "image_cifar10_plain"  # Reuse CIFAR-10 plain data.
+
+  def preprocess_example(self, example, mode, unused_hparams):
+    example["inputs"] = resize_by_area(example["inputs"], 8)
+    return example
+
+
 @registry.register_problem
 class Img2imgCifar10(ImageCifar10):
   """CIFAR-10 rescaled to 8x8 for input and 32x32 for output."""
@@ -663,14 +673,10 @@ def dataset_filename(self):
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
 
-    def resize(img, size):
-      return tf.to_int64(
-          tf.image.resize_images(img, [size, size], tf.image.ResizeMethod.AREA))
-
     inputs = example["inputs"]
     # For Img2Img resize input and output images as desired.
-    example["inputs"] = resize(inputs, 8)
-    example["targets"] = resize(inputs, 32)
+    example["inputs"] = resize_by_area(inputs, 8)
+    example["targets"] = resize_by_area(inputs, 32)
     return example
 
   def hparams(self, defaults, unused_model_hparams):

tensor2tensor/layers/common_hparams.py

@@ -62,6 +62,7 @@ def basic_params1():
       learning_rate_cosine_cycle_steps=250000,
       learning_rate=0.1,
       sampling_method="argmax",  # "argmax" or "random"
+      sampling_temp=1.0,  # temperature for sampling
       problem_choice="adaptive",  # "uniform", "adaptive", "distributed"
       # expand the logits a piece at a time - saves memory.
       factored_logits=int(False),

tensor2tensor/utils/t2t_model.py

@@ -427,15 +427,17 @@ def sample(self, features, last_position_only=False):
     else:
       assert self._hparams.sampling_method == "random"
 
-      def _multinomial_squeeze(logits):
-        reshaped_logits = tf.reshape(logits, [-1, tf.shape(logits)[-1]])
+      def _multinomial_squeeze(logits, temperature=1.0):
+        reshaped_logits = (
+            tf.reshape(logits, [-1, tf.shape(logits)[-1]])/temperature)
         choices = tf.multinomial(reshaped_logits, 1)
         choices = tf.reshape(choices,
                              tf.shape(logits)[:logits.get_shape().ndims - 1])
         return choices
 
       sharded_samples = self._data_parallelism(_multinomial_squeeze,
-                                               sharded_logits)
+                                               sharded_logits,
+                                               self._hparams.sampling_temp)
     return tf.concat(sharded_samples, 0), sharded_logits, losses
 
   def _shard_features(self, features):  # pylint: disable=missing-docstring