Project import generated by Copybara.

PiperOrigin-RevId: 176534973

research/gan/mnist/eval.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Evaluates a TFGAN trained MNIST model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+
+import tensorflow as tf
+
+import data_provider
+import networks
+import util
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+tfgan = tf.contrib.gan
+
+
+flags.DEFINE_string('checkpoint_dir', '/tmp/mnist/',
+                    'Directory where the model was written to.')
+
+flags.DEFINE_string('eval_dir', '/tmp/mnist/',
+                    'Directory where the results are saved to.')
+
+flags.DEFINE_string('dataset_dir', None, 'Location of data.')
+
+flags.DEFINE_integer('num_images_generated', 1000,
+                     'Number of images to generate at once.')
+
+flags.DEFINE_boolean('eval_real_images', False,
+                     'If `True`, run Inception network on real images.')
+
+flags.DEFINE_integer('noise_dims', 64,
+                     'Dimensions of the generator noise vector')
+
+flags.DEFINE_string('classifier_filename', None,
+                    'Location of the pretrained classifier. If `None`, use '
+                    'default.')
+
+flags.DEFINE_integer('max_number_of_evaluations', None,
+                     'Number of times to run evaluation. If `None`, run '
+                     'forever.')
+
+
+def main(_, run_eval_loop=True):
+  # Fetch real images.
+  with tf.name_scope('inputs'):
+    real_images, _, _ = data_provider.provide_data(
+        'train', FLAGS.num_images_generated, FLAGS.dataset_dir)
+
+  image_write_ops = None
+  if FLAGS.eval_real_images:
+    tf.summary.scalar('MNIST_Classifier_score',
+                      util.mnist_score(real_images, FLAGS.classifier_filename))
+  else:
+    # In order for variables to load, use the same variable scope as in the
+    # train job.
+    with tf.variable_scope('Generator'):
+      images = networks.unconditional_generator(
+          tf.random_normal([FLAGS.num_images_generated, FLAGS.noise_dims]))
+    tf.summary.scalar('MNIST_Frechet_distance',
+                      util.mnist_frechet_distance(
+                          real_images, images, FLAGS.classifier_filename))
+    tf.summary.scalar('MNIST_Classifier_score',
+                      util.mnist_score(images, FLAGS.classifier_filename))
+    if FLAGS.num_images_generated >= 100:
+      reshaped_images = tfgan.eval.image_reshaper(
+          images[:100, ...], num_cols=10)
+      uint8_images = data_provider.float_image_to_uint8(reshaped_images)
+      image_write_ops = tf.write_file(
+          '%s/%s'% (FLAGS.eval_dir, 'unconditional_gan.png'),
+          tf.image.encode_png(uint8_images[0]))
+
+  # For unit testing, use `run_eval_loop=False`.
+  if not run_eval_loop: return
+  tf.contrib.training.evaluate_repeatedly(
+      FLAGS.checkpoint_dir,
+      hooks=[tf.contrib.training.SummaryAtEndHook(FLAGS.eval_dir),
+             tf.contrib.training.StopAfterNEvalsHook(1)],
+      eval_ops=image_write_ops,
+      max_number_of_evaluations=FLAGS.max_number_of_evaluations)
+
+
+if __name__ == '__main__':
+  tf.app.run()

research/gan/mnist/eval_test.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for tfgan.examples.mnist.eval."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+
+import tensorflow as tf
+from google3.third_party.tensorflow_models.gan.mnist  import eval  # pylint:disable=redefined-builtin
+
+
+class EvalTest(tf.test.TestCase):
+
+  def _test_build_graph_helper(self, eval_real_images):
+    tf.flags.FLAGS.eval_real_images = eval_real_images
+    eval.main(None, run_eval_loop=False)
+
+  def test_build_graph_realdata(self):
+    self._test_build_graph_helper(True)
+
+  def test_build_graph_generateddata(self):
+    self._test_build_graph_helper(False)
+
+if __name__ == '__main__':
+  tf.test.main()

research/gan/mnist/infogan_eval.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Evaluates an InfoGAN TFGAN trained MNIST model.
+
+The image visualizations, as in https://arxiv.org/abs/1606.03657, show the
+effect of varying a specific latent variable on the image. Each visualization
+focuses on one of the three structured variables. Columns have two of the three
+variables fixed, while the third one is varied. Different rows have different
+random samples from the remaining latents.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import numpy as np
+
+import tensorflow as tf
+
+import data_provider
+import networks
+import util
+
+flags = tf.flags
+tfgan = tf.contrib.gan
+
+
+flags.DEFINE_string('checkpoint_dir', '/tmp/mnist/',
+                    'Directory where the model was written to.')
+
+flags.DEFINE_string('eval_dir', '/tmp/mnist/',
+                    'Directory where the results are saved to.')
+
+flags.DEFINE_integer('noise_samples', 6,
+                     'Number of samples to draw from the continuous structured '
+                     'noise.')
+
+flags.DEFINE_integer('unstructured_noise_dims', 62,
+                     'The number of dimensions of the unstructured noise.')
+
+flags.DEFINE_integer('continuous_noise_dims', 2,
+                     'The number of dimensions of the continuous noise.')
+
+flags.DEFINE_string('classifier_filename', None,
+                    'Location of the pretrained classifier. If `None`, use '
+                    'default.')
+
+flags.DEFINE_integer('max_number_of_evaluations', None,
+                     'Number of times to run evaluation. If `None`, run '
+                     'forever.')
+
+CAT_SAMPLE_POINTS = np.arange(0, 10)
+CONT_SAMPLE_POINTS = np.linspace(-2.0, 2.0, 10)
+FLAGS = flags.FLAGS
+
+
+def main(_, run_eval_loop=True):
+  with tf.name_scope('inputs'):
+    noise_args = (FLAGS.noise_samples, CAT_SAMPLE_POINTS, CONT_SAMPLE_POINTS,
+                  FLAGS.unstructured_noise_dims, FLAGS.continuous_noise_dims)
+    # Use fixed noise vectors to illustrate the effect of each dimension.
+    display_noise1 = util.get_eval_noise_categorical(*noise_args)
+    display_noise2 = util.get_eval_noise_continuous_dim1(*noise_args)
+    display_noise3 = util.get_eval_noise_continuous_dim2(*noise_args)
+    _validate_noises([display_noise1, display_noise2, display_noise3])
+
+  # Visualize the effect of each structured noise dimension on the generated
+  # image.
+  generator_fn = lambda x: networks.infogan_generator(x, len(CAT_SAMPLE_POINTS))
+  with tf.variable_scope('Generator') as genscope:  # Same scope as in training.
+    categorical_images = generator_fn(display_noise1)
+  reshaped_categorical_img = tfgan.eval.image_reshaper(
+      categorical_images, num_cols=len(CAT_SAMPLE_POINTS))
+  tf.summary.image('categorical', reshaped_categorical_img, max_outputs=1)
+
+  with tf.variable_scope(genscope, reuse=True):
+    continuous1_images = generator_fn(display_noise2)
+  reshaped_continuous1_img = tfgan.eval.image_reshaper(
+      continuous1_images, num_cols=len(CONT_SAMPLE_POINTS))
+  tf.summary.image('continuous1', reshaped_continuous1_img, max_outputs=1)
+
+  with tf.variable_scope(genscope, reuse=True):
+    continuous2_images = generator_fn(display_noise3)
+  reshaped_continuous2_img = tfgan.eval.image_reshaper(
+      continuous2_images, num_cols=len(CONT_SAMPLE_POINTS))
+  tf.summary.image('continuous2', reshaped_continuous2_img, max_outputs=1)
+
+  # Evaluate image quality.
+  all_images = tf.concat(
+      [categorical_images, continuous1_images, continuous2_images], 0)
+  tf.summary.scalar('MNIST_Classifier_score',
+                    util.mnist_score(all_images, FLAGS.classifier_filename))
+
+  # Write images to disk.
+  image_write_ops = []
+  image_write_ops.append(_get_write_image_ops(
+      FLAGS.eval_dir, 'categorical_infogan.png', reshaped_categorical_img[0]))
+  image_write_ops.append(_get_write_image_ops(
+      FLAGS.eval_dir, 'continuous1_infogan.png', reshaped_continuous1_img[0]))
+  image_write_ops.append(_get_write_image_ops(
+      FLAGS.eval_dir, 'continuous2_infogan.png', reshaped_continuous2_img[0]))
+
+  # For unit testing, use `run_eval_loop=False`.
+  if not run_eval_loop: return
+  tf.contrib.training.evaluate_repeatedly(
+      FLAGS.checkpoint_dir,
+      hooks=[tf.contrib.training.SummaryAtEndHook(FLAGS.eval_dir),
+             tf.contrib.training.StopAfterNEvalsHook(1)],
+      eval_ops=image_write_ops,
+      max_number_of_evaluations=FLAGS.max_number_of_evaluations)
+
+
+def _validate_noises(noises):
+  """Sanity check on constructed noise tensors.
+
+  Args:
+    noises: List of 3-tuples of noise vectors.
+  """
+  assert isinstance(noises, (list, tuple))
+  for noise_l in noises:
+    assert len(noise_l) == 3
+    assert isinstance(noise_l[0], np.ndarray)
+    batch_dim = noise_l[0].shape[0]
+    for i, noise in enumerate(noise_l):
+      assert isinstance(noise, np.ndarray)
+      # Check that batch dimensions are all the same.
+      assert noise.shape[0] == batch_dim
+
+      # Check that shapes for corresponding noises are the same.
+      assert noise.shape == noises[0][i].shape
+
+
+def _get_write_image_ops(eval_dir, filename, images):
+  """Create Ops that write images to disk."""
+  return tf.write_file(
+      '%s/%s'% (eval_dir, filename),
+      tf.image.encode_png(data_provider.float_image_to_uint8(images)))
+
+
+if __name__ == '__main__':
+  tf.app.run()

research/gan/mnist/infogan_eval_test.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for tfgan.examples.mnist.infogan_eval."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import tensorflow as tf
+import infogan_eval
+
+
+class MnistInfoGANEvalTest(tf.test.TestCase):
+
+  def test_build_graph(self):
+    infogan_eval.main(None, run_eval_loop=False)
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/gan/mnist/launch_jobs.sh

+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+#!/bin/bash
+#
+# This script performs the following operations:
+# 1. Downloads the MNIST dataset.
+# 2. Trains an unconditional, conditional, or InfoGAN model on the MNIST
+#    training set.
+# 3. Evaluates the models and writes sample images to disk.
+#
+# These examples are intended to be fast. For better final results, tune
+# hyperparameters or train longer.
+#
+# NOTE: Each training step takes about 0.5 second with a batch size of 32 on
+# CPU. On GPU, it takes ~5 milliseconds.
+#
+# With the default batch size and number of steps, train times are:
+#
+#   unconditional: CPU: 800  steps, ~10 min   GPU: 800  steps, ~1 min
+#   conditional:   CPU: 2000 steps, ~20 min   GPU: 2000 steps, ~2 min
+#   infogan:       CPU: 3000 steps, ~20 min   GPU: 3000 steps, ~6 min
+#
+# Usage:
+# cd models/research/gan/mnist
+# ./launch_jobs.sh ${gan_type} ${git_repo}
+set -e
+
+# Type of GAN to run. Right now, options are `unconditional`, `conditional`, or
+# `infogan`.
+gan_type=$1
+if ! [[ "$gan_type" =~ ^(unconditional|conditional|infogan) ]]; then
+  echo "'gan_type' must be one of: 'unconditional', 'conditional', 'infogan'."
+  exit
+fi
+
+# Location of the git repository.
+git_repo=$2
+if [[ "$git_repo" == "" ]]; then
+  echo "'git_repo' must not be empty."
+  exit
+fi
+
+# Base name for where the checkpoint and logs will be saved to.
+TRAIN_DIR=/tmp/mnist-model
+
+# Base name for where the evaluation images will be saved to.
+EVAL_DIR=/tmp/mnist-model/eval
+
+# Where the dataset is saved to.
+DATASET_DIR=/tmp/mnist-data
+
+# Location of the classifier frozen graph used for evaluation.
+FROZEN_GRAPH="${git_repo}/research/gan/mnist/data/classify_mnist_graph_def.pb"
+
+export PYTHONPATH=$PYTHONPATH:$git_repo:$git_repo/research:$git_repo/research/slim
+
+# A helper function for printing pretty output.
+Banner () {
+  local text=$1
+  local green='\033[0;32m'
+  local nc='\033[0m'  # No color.
+  echo -e "${green}${text}${nc}"
+}
+
+# Download the dataset.
+python "${git_repo}/research/slim/download_and_convert_data.py" \
+  --dataset_name=mnist \
+  --dataset_dir=${DATASET_DIR}
+
+# Run unconditional GAN.
+if [[ "$gan_type" == "unconditional" ]]; then
+  UNCONDITIONAL_TRAIN_DIR="${TRAIN_DIR}/unconditional"
+  UNCONDITIONAL_EVAL_DIR="${EVAL_DIR}/unconditional"
+  NUM_STEPS=3000
+  # Run training.
+  Banner "Starting training unconditional GAN for ${NUM_STEPS} steps..."
+  python "${git_repo}/research/gan/mnist/train.py" \
+    --train_log_dir=${UNCONDITIONAL_TRAIN_DIR} \
+    --dataset_dir=${DATASET_DIR} \
+    --max_number_of_steps=${NUM_STEPS} \
+    --gan_type="unconditional" \
+    --alsologtostderr
+  Banner "Finished training unconditional GAN ${NUM_STEPS} steps."
+
+  # Run evaluation.
+  Banner "Starting evaluation of unconditional GAN..."
+  python "${git_repo}/research/gan/mnist/eval.py" \
+    --checkpoint_dir=${UNCONDITIONAL_TRAIN_DIR} \
+    --eval_dir=${UNCONDITIONAL_EVAL_DIR} \
+    --dataset_dir=${DATASET_DIR} \
+    --eval_real_images=false \
+    --classifier_filename=${FROZEN_GRAPH} \
+    --max_number_of_evaluation=1
+  Banner "Finished unconditional evaluation. See ${UNCONDITIONAL_EVAL_DIR} for output images."
+fi
+
+# Run conditional GAN.
+if [[ "$gan_type" == "conditional" ]]; then
+  CONDITIONAL_TRAIN_DIR="${TRAIN_DIR}/conditional"
+  CONDITIONAL_EVAL_DIR="${EVAL_DIR}/conditional"
+  NUM_STEPS=3000
+  # Run training.
+  Banner "Starting training conditional GAN for ${NUM_STEPS} steps..."
+  python "${git_repo}/research/gan/mnist/train.py" \
+    --train_log_dir=${CONDITIONAL_TRAIN_DIR} \
+    --dataset_dir=${DATASET_DIR} \
+    --max_number_of_steps=${NUM_STEPS} \
+    --gan_type="conditional" \
+    --alsologtostderr
+  Banner "Finished training conditional GAN ${NUM_STEPS} steps."
+
+  # Run evaluation.
+  Banner "Starting evaluation of conditional GAN..."
+  python "${git_repo}/research/gan/mnist/conditional_eval.py" \
+    --checkpoint_dir=${CONDITIONAL_TRAIN_DIR} \
+    --eval_dir=${CONDITIONAL_EVAL_DIR} \
+    --classifier_filename=${FROZEN_GRAPH} \
+    --max_number_of_evaluation=1
+  Banner "Finished conditional evaluation. See ${CONDITIONAL_EVAL_DIR} for output images."
+fi
+
+# Run InfoGAN.
+if [[ "$gan_type" == "infogan" ]]; then
+  INFOGAN_TRAIN_DIR="${TRAIN_DIR}/infogan"
+  INFOGAN_EVAL_DIR="${EVAL_DIR}/infogan"
+  NUM_STEPS=3000
+  # Run training.
+  Banner "Starting training infogan GAN for ${NUM_STEPS} steps..."
+  python "${git_repo}/research/gan/mnist/train.py" \
+    --train_log_dir=${INFOGAN_TRAIN_DIR} \
+    --dataset_dir=${DATASET_DIR} \
+    --max_number_of_steps=${NUM_STEPS} \
+    --gan_type="infogan" \
+    --alsologtostderr
+  Banner "Finished training InfoGAN ${NUM_STEPS} steps."
+
+  # Run evaluation.
+  Banner "Starting evaluation of infogan..."
+  python "${git_repo}/research/gan/mnist/infogan_eval.py" \
+    --checkpoint_dir=${INFOGAN_TRAIN_DIR} \
+    --eval_dir=${INFOGAN_EVAL_DIR} \
+    --classifier_filename=${FROZEN_GRAPH} \
+    --max_number_of_evaluation=1
+  Banner "Finished InfoGAN evaluation. See ${INFOGAN_EVAL_DIR} for output images."
+fi

research/gan/mnist/networks.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Networks for MNIST example using TFGAN."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+ds = tf.contrib.distributions
+layers = tf.contrib.layers
+tfgan = tf.contrib.gan
+
+
+def _generator_helper(
+    noise, is_conditional, one_hot_labels, weight_decay):
+  """Core MNIST generator.
+
+  This function is reused between the different GAN modes (unconditional,
+  conditional, etc).
+
+  Args:
+    noise: A 2D Tensor of shape [batch size, noise dim].
+    is_conditional: Whether to condition on labels.
+    one_hot_labels: Optional labels for conditioning.
+    weight_decay: The value of the l2 weight decay.
+
+  Returns:
+    A generated image in the range [-1, 1].
+  """
+  with tf.contrib.framework.arg_scope(
+      [layers.fully_connected, layers.conv2d_transpose],
+      activation_fn=tf.nn.relu, normalizer_fn=layers.batch_norm,
+      weights_regularizer=layers.l2_regularizer(weight_decay)):
+    net = layers.fully_connected(noise, 1024)
+    if is_conditional:
+      net = tfgan.features.condition_tensor_from_onehot(net, one_hot_labels)
+    net = layers.fully_connected(net, 7 * 7 * 128)
+    net = tf.reshape(net, [-1, 7, 7, 128])
+    net = layers.conv2d_transpose(net, 64, [4, 4], stride=2)
+    net = layers.conv2d_transpose(net, 32, [4, 4], stride=2)
+    # Make sure that generator output is in the same range as `inputs`
+    # ie [-1, 1].
+    net = layers.conv2d(
+        net, 1, [4, 4], normalizer_fn=None, activation_fn=tf.tanh)
+
+    return net
+
+
+def unconditional_generator(noise, weight_decay=2.5e-5):
+  """Generator to produce unconditional MNIST images.
+
+  Args:
+    noise: A single Tensor representing noise.
+    weight_decay: The value of the l2 weight decay.
+
+  Returns:
+    A generated image in the range [-1, 1].
+  """
+  return _generator_helper(noise, False, None, weight_decay)
+
+
+def conditional_generator(inputs, weight_decay=2.5e-5):
+  """Generator to produce MNIST images conditioned on class.
+
+  Args:
+    inputs: A 2-tuple of Tensors (noise, one_hot_labels).
+    weight_decay: The value of the l2 weight decay.
+
+  Returns:
+    A generated image in the range [-1, 1].
+  """
+  noise, one_hot_labels = inputs
+  return _generator_helper(noise, True, one_hot_labels, weight_decay)
+
+
+def infogan_generator(inputs, categorical_dim, weight_decay=2.5e-5):
+  """InfoGAN generator network on MNIST digits.
+
+  Based on a paper https://arxiv.org/abs/1606.03657, their code
+  https://github.com/openai/InfoGAN, and code by pooleb@.
+
+  Args:
+    inputs: A 3-tuple of Tensors (unstructured_noise, categorical structured
+      noise, continuous structured noise). `inputs[0]` and `inputs[2]` must be
+      2D, and `inputs[1]` must be 1D. All must have the same first dimension.
+    categorical_dim: Dimensions of the incompressible categorical noise.
+    weight_decay: The value of the l2 weight decay.
+
+  Returns:
+    A generated image in the range [-1, 1].
+  """
+  unstructured_noise, cat_noise, cont_noise = inputs
+  cat_noise_onehot = tf.one_hot(cat_noise, categorical_dim)
+  all_noise = tf.concat(
+      [unstructured_noise, cat_noise_onehot, cont_noise], axis=1)
+  return _generator_helper(all_noise, False, None, weight_decay)
+
+
+_leaky_relu = lambda x: tf.nn.leaky_relu(x, alpha=0.01)
+
+
+def _discriminator_helper(img, is_conditional, one_hot_labels, weight_decay):
+  """Core MNIST discriminator.
+
+  This function is reused between the different GAN modes (unconditional,
+  conditional, etc).
+
+  Args:
+    img: Real or generated MNIST digits. Should be in the range [-1, 1].
+    is_conditional: Whether to condition on labels.
+    one_hot_labels: Labels to optionally condition the network on.
+    weight_decay: The L2 weight decay.
+
+  Returns:
+    Final fully connected discriminator layer. [batch_size, 1024].
+  """
+  with tf.contrib.framework.arg_scope(
+      [layers.conv2d, layers.fully_connected],
+      activation_fn=_leaky_relu, normalizer_fn=None,
+      weights_regularizer=layers.l2_regularizer(weight_decay),
+      biases_regularizer=layers.l2_regularizer(weight_decay)):
+    net = layers.conv2d(img, 64, [4, 4], stride=2)
+    net = layers.conv2d(net, 128, [4, 4], stride=2)
+    net = layers.flatten(net)
+    if is_conditional:
+      net = tfgan.features.condition_tensor_from_onehot(net, one_hot_labels)
+    net = layers.fully_connected(net, 1024, normalizer_fn=layers.layer_norm)
+
+    return net
+
+
+def unconditional_discriminator(img, unused_conditioning, weight_decay=2.5e-5):
+  """Discriminator network on unconditional MNIST digits.
+
+  Args:
+    img: Real or generated MNIST digits. Should be in the range [-1, 1].
+    unused_conditioning: The TFGAN API can help with conditional GANs, which
+      would require extra `condition` information to both the generator and the
+      discriminator. Since this example is not conditional, we do not use this
+      argument.
+    weight_decay: The L2 weight decay.
+
+  Returns:
+    Logits for the probability that the image is real.
+  """
+  net = _discriminator_helper(img, False, None, weight_decay)
+  return layers.linear(net, 1)
+
+
+def conditional_discriminator(img, conditioning, weight_decay=2.5e-5):
+  """Conditional discriminator network on MNIST digits.
+
+  Args:
+    img: Real or generated MNIST digits. Should be in the range [-1, 1].
+    conditioning: A 2-tuple of Tensors representing (noise, one_hot_labels).
+    weight_decay: The L2 weight decay.
+
+  Returns:
+    Logits for the probability that the image is real.
+  """
+  _, one_hot_labels = conditioning
+  net = _discriminator_helper(img, True, one_hot_labels, weight_decay)
+  return layers.linear(net, 1)
+
+
+def infogan_discriminator(img, unused_conditioning, weight_decay=2.5e-5,
+                          categorical_dim=10, continuous_dim=2):
+  """InfoGAN discriminator network on MNIST digits.
+
+  Based on a paper https://arxiv.org/abs/1606.03657, their code
+  https://github.com/openai/InfoGAN, and code by pooleb@.
+
+  Args:
+    img: Real or generated MNIST digits. Should be in the range [-1, 1].
+    unused_conditioning: The TFGAN API can help with conditional GANs, which
+      would require extra `condition` information to both the generator and the
+      discriminator. Since this example is not conditional, we do not use this
+      argument.
+    weight_decay: The L2 weight decay.
+    categorical_dim: Dimensions of the incompressible categorical noise.
+    continuous_dim: Dimensions of the incompressible continuous noise.
+
+  Returns:
+    Logits for the probability that the image is real, and a list of posterior
+    distributions for each of the noise vectors.
+  """
+  net = _discriminator_helper(img, False, None, weight_decay)
+  logits_real = layers.fully_connected(net, 1, activation_fn=None)
+
+  # Recognition network for latent variables has an additional layer
+  with tf.contrib.framework.arg_scope([layers.batch_norm], is_training=False):
+    encoder = layers.fully_connected(net, 128, normalizer_fn=layers.batch_norm,
+                                     activation_fn=_leaky_relu)
+
+  # Compute logits for each category of categorical latent.
+  logits_cat = layers.fully_connected(
+      encoder, categorical_dim, activation_fn=None)
+  q_cat = ds.Categorical(logits_cat)
+
+  # Compute mean for Gaussian posterior of continuous latents.
+  mu_cont = layers.fully_connected(encoder, continuous_dim, activation_fn=None)
+  sigma_cont = tf.ones_like(mu_cont)
+  q_cont = ds.Normal(loc=mu_cont, scale=sigma_cont)
+
+  return logits_real, [q_cat, q_cont]

research/gan/mnist/train.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Trains a generator on MNIST data."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+
+import tensorflow as tf
+
+import data_provider
+import networks
+import util
+
+flags = tf.flags
+tfgan = tf.contrib.gan
+
+
+flags.DEFINE_integer('batch_size', 32, 'The number of images in each batch.')
+
+flags.DEFINE_string('train_log_dir', '/tmp/mnist/',
+                    'Directory where to write event logs.')
+
+flags.DEFINE_string('dataset_dir', None, 'Location of data.')
+
+flags.DEFINE_integer('max_number_of_steps', 20000,
+                     'The maximum number of gradient steps.')
+
+flags.DEFINE_string(
+    'gan_type', 'unconditional',
+    'Either `unconditional`, `conditional`, or `infogan`.')
+
+flags.DEFINE_integer(
+    'grid_size', 5, 'Grid size for image visualization.')
+
+
+flags.DEFINE_integer(
+    'noise_dims', 64, 'Dimensions of the generator noise vector.')
+
+FLAGS = flags.FLAGS
+
+
+def _learning_rate(gan_type):
+  # First is generator learning rate, second is discriminator learning rate.
+  return {
+      'unconditional': (1e-3, 1e-4),
+      'conditional': (1e-5, 1e-4),
+      'infogan': (0.001, 9e-5),
+  }[gan_type]
+
+
+def main(_):
+  if not tf.gfile.Exists(FLAGS.train_log_dir):
+    tf.gfile.MakeDirs(FLAGS.train_log_dir)
+
+  # Force all input processing onto CPU in order to reserve the GPU for
+  # the forward inference and back-propagation.
+  with tf.name_scope('inputs'):
+    with tf.device('/cpu:0'):
+      images, one_hot_labels, _ = data_provider.provide_data(
+          'train', FLAGS.batch_size, FLAGS.dataset_dir, num_threads=4)
+
+  # Define the GANModel tuple. Optionally, condition the GAN on the label or
+  # use an InfoGAN to learn a latent representation.
+  if FLAGS.gan_type == 'unconditional':
+    gan_model = tfgan.gan_model(
+        generator_fn=networks.unconditional_generator,
+        discriminator_fn=networks.unconditional_discriminator,
+        real_data=images,
+        generator_inputs=tf.random_normal(
+            [FLAGS.batch_size, FLAGS.noise_dims]))
+  elif FLAGS.gan_type == 'conditional':
+    noise = tf.random_normal([FLAGS.batch_size, FLAGS.noise_dims])
+    gan_model = tfgan.gan_model(
+        generator_fn=networks.conditional_generator,
+        discriminator_fn=networks.conditional_discriminator,
+        real_data=images,
+        generator_inputs=(noise, one_hot_labels))
+  elif FLAGS.gan_type == 'infogan':
+    cat_dim, cont_dim = 10, 2
+    generator_fn = functools.partial(
+        networks.infogan_generator, categorical_dim=cat_dim)
+    discriminator_fn = functools.partial(
+        networks.infogan_discriminator, categorical_dim=cat_dim,
+        continuous_dim=cont_dim)
+    unstructured_inputs, structured_inputs = util.get_infogan_noise(
+        FLAGS.batch_size, cat_dim, cont_dim, FLAGS.noise_dims)
+    gan_model = tfgan.infogan_model(
+        generator_fn=generator_fn,
+        discriminator_fn=discriminator_fn,
+        real_data=images,
+        unstructured_generator_inputs=unstructured_inputs,
+        structured_generator_inputs=structured_inputs)
+  tfgan.eval.add_gan_model_image_summaries(gan_model, FLAGS.grid_size)
+
+  # Get the GANLoss tuple. You can pass a custom function, use one of the
+  # already-implemented losses from the losses library, or use the defaults.
+  with tf.name_scope('loss'):
+    mutual_information_penalty_weight = (1.0 if FLAGS.gan_type == 'infogan'
+                                         else 0.0)
+    gan_loss = tfgan.gan_loss(
+        gan_model,
+        gradient_penalty_weight=1.0,
+        mutual_information_penalty_weight=mutual_information_penalty_weight,
+        add_summaries=True)
+    tfgan.eval.add_regularization_loss_summaries(gan_model)
+
+  # Get the GANTrain ops using custom optimizers.
+  with tf.name_scope('train'):
+    gen_lr, dis_lr = _learning_rate(FLAGS.gan_type)
+    train_ops = tfgan.gan_train_ops(
+        gan_model,
+        gan_loss,
+        generator_optimizer=tf.train.AdamOptimizer(gen_lr, 0.5),
+        discriminator_optimizer=tf.train.AdamOptimizer(dis_lr, 0.5),
+        summarize_gradients=True,
+        aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
+
+  # Run the alternating training loop. Skip it if no steps should be taken
+  # (used for graph construction tests).
+  status_message = tf.string_join(
+      ['Starting train step: ',
+       tf.as_string(tf.train.get_or_create_global_step())],
+      name='status_message')
+  if FLAGS.max_number_of_steps == 0: return
+  tfgan.gan_train(
+      train_ops,
+      hooks=[tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
+             tf.train.LoggingTensorHook([status_message], every_n_iter=10)],
+      logdir=FLAGS.train_log_dir,
+      get_hooks_fn=tfgan.get_joint_train_hooks())
+
+if __name__ == '__main__':
+  tf.logging.set_verbosity(tf.logging.INFO)
+  tf.app.run()

research/gan/mnist/train_test.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for tfgan.examples.mnist.train."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import numpy as np
+
+import tensorflow as tf
+import train
+
+FLAGS = tf.flags.FLAGS
+mock = tf.test.mock
+
+
+class TrainTest(tf.test.TestCase):
+
+  def test_run_one_train_step(self):
+    FLAGS.max_number_of_steps = 1
+    FLAGS.gan_type = 'unconditional'
+    FLAGS.batch_size = 5
+    FLAGS.grid_size = 1
+    tf.set_random_seed(1234)
+
+    # Mock input pipeline.
+    mock_imgs = np.zeros([FLAGS.batch_size, 28, 28, 1], dtype=np.float32)
+    mock_lbls = np.concatenate(
+        (np.ones([FLAGS.batch_size, 1], dtype=np.int32),
+         np.zeros([FLAGS.batch_size, 9], dtype=np.int32)), axis=1)
+    with mock.patch.object(train, 'data_provider') as mock_data_provider:
+      mock_data_provider.provide_data.return_value = (
+          mock_imgs, mock_lbls, None)
+      train.main(None)
+
+  def _test_build_graph_helper(self, gan_type):
+    FLAGS.max_number_of_steps = 0
+    FLAGS.gan_type = gan_type
+
+    # Mock input pipeline.
+    mock_imgs = np.zeros([FLAGS.batch_size, 28, 28, 1], dtype=np.float32)
+    mock_lbls = np.concatenate(
+        (np.ones([FLAGS.batch_size, 1], dtype=np.int32),
+         np.zeros([FLAGS.batch_size, 9], dtype=np.int32)), axis=1)
+    with mock.patch.object(train, 'data_provider') as mock_data_provider:
+      mock_data_provider.provide_data.return_value = (
+          mock_imgs, mock_lbls, None)
+      train.main(None)
+
+  def test_build_graph_unconditional(self):
+    self._test_build_graph_helper('unconditional')
+
+  def test_build_graph_conditional(self):
+    self._test_build_graph_helper('conditional')
+
+  def test_build_graph_infogan(self):
+    self._test_build_graph_helper('infogan')
+
+if __name__ == '__main__':
+  tf.test.main()

research/gan/mnist/util.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""A utility for evaluating MNIST generative models.
+
+These functions use a pretrained MNIST classifier with ~99% eval accuracy to
+measure various aspects of the quality of generated MNIST digits.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import numpy as np
+import tensorflow as tf
+
+ds = tf.contrib.distributions
+tfgan = tf.contrib.gan
+
+
+__all__ = [
+    'mnist_score',
+    'mnist_frechet_distance',
+    'mnist_cross_entropy',
+    'get_eval_noise_categorical',
+    'get_eval_noise_continuous_dim1',
+    'get_eval_noise_continuous_dim2',
+    'get_infogan_noise',
+]
+
+
+# Prepend `../`, since paths start from `third_party/tensorflow`.
+MODEL_GRAPH_DEF = '../tensorflow_models/gan/mnist/data/classify_mnist_graph_def.pb'
+INPUT_TENSOR = 'inputs:0'
+OUTPUT_TENSOR = 'logits:0'
+
+
+def mnist_score(images, graph_def_filename=None, input_tensor=INPUT_TENSOR,
+                output_tensor=OUTPUT_TENSOR, num_batches=1):
+  """Get MNIST logits of a fully-trained classifier.
+
+  Args:
+    images: A minibatch tensor of MNIST digits. Shape must be
+      [batch, 28, 28, 1].
+    graph_def_filename: Location of a frozen GraphDef binary file on disk. If
+      `None`, uses a default graph.
+    input_tensor: GraphDef's input tensor name.
+    output_tensor: GraphDef's output tensor name.
+    num_batches: Number of batches to split `generated_images` in to in order to
+      efficiently run them through Inception.
+
+  Returns:
+    A logits tensor of [batch, 10].
+  """
+  images.shape.assert_is_compatible_with([None, 28, 28, 1])
+
+  graph_def = _graph_def_from_par_or_disk(graph_def_filename)
+  mnist_classifier_fn = lambda x: tfgan.eval.run_image_classifier(  # pylint: disable=g-long-lambda
+      x, graph_def, input_tensor, output_tensor)
+
+  score = tfgan.eval.classifier_score(
+      images, mnist_classifier_fn, num_batches)
+  score.shape.assert_is_compatible_with([])
+
+  return score
+
+
+def mnist_frechet_distance(real_images, generated_images,
+                           graph_def_filename=None, input_tensor=INPUT_TENSOR,
+                           output_tensor=OUTPUT_TENSOR, num_batches=1):
+  """Frechet distance between real and generated images.
+
+  This technique is described in detail in https://arxiv.org/abs/1706.08500.
+  Please see TFGAN for implementation details
+  (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/gan/
+  python/eval/python/classifier_metrics_impl.py).
+
+  Args:
+    real_images: Real images to use to compute Frechet Inception distance.
+    generated_images: Generated images to use to compute Frechet Inception
+      distance.
+    graph_def_filename: Location of a frozen GraphDef binary file on disk. If
+      `None`, uses a default graph.
+    input_tensor: GraphDef's input tensor name.
+    output_tensor: GraphDef's output tensor name.
+    num_batches: Number of batches to split images into in order to
+      efficiently run them through the classifier network.
+
+  Returns:
+    The Frechet distance. A floating-point scalar.
+  """
+  real_images.shape.assert_is_compatible_with([None, 28, 28, 1])
+  generated_images.shape.assert_is_compatible_with([None, 28, 28, 1])
+
+  graph_def = _graph_def_from_par_or_disk(graph_def_filename)
+  mnist_classifier_fn = lambda x: tfgan.eval.run_image_classifier(  # pylint: disable=g-long-lambda
+      x, graph_def, input_tensor, output_tensor)
+
+  frechet_distance = tfgan.eval.frechet_classifier_distance(
+      real_images, generated_images, mnist_classifier_fn, num_batches)
+  frechet_distance.shape.assert_is_compatible_with([])
+
+  return frechet_distance
+
+
+def mnist_cross_entropy(images, one_hot_labels, graph_def_filename=None,
+                        input_tensor=INPUT_TENSOR, output_tensor=OUTPUT_TENSOR):
+  """Returns the cross entropy loss of the classifier on images.
+
+  Args:
+    images: A minibatch tensor of MNIST digits. Shape must be
+      [batch, 28, 28, 1].
+    one_hot_labels: The one hot label of the examples. Tensor size is
+      [batch, 10].
+    graph_def_filename: Location of a frozen GraphDef binary file on disk. If
+      `None`, uses a default graph embedded in the par file.
+    input_tensor: GraphDef's input tensor name.
+    output_tensor: GraphDef's output tensor name.
+
+  Returns:
+    A scalar Tensor representing the cross entropy of the image minibatch.
+  """
+  graph_def = _graph_def_from_par_or_disk(graph_def_filename)
+
+  logits = tfgan.eval.run_image_classifier(
+      images, graph_def, input_tensor, output_tensor)
+  return tf.losses.softmax_cross_entropy(
+      one_hot_labels, logits, loss_collection=None)
+
+
+# (joelshor): Refactor the `eval_noise` functions to reuse code.
+def get_eval_noise_categorical(
+    noise_samples, categorical_sample_points, continuous_sample_points,
+    unstructured_noise_dims, continuous_noise_dims):
+  """Create noise showing impact of categorical noise in InfoGAN.
+
+  Categorical noise is constant across columns. Other noise is constant across
+  rows.
+
+  Args:
+    noise_samples: Number of non-categorical noise samples to use.
+    categorical_sample_points: Possible categorical noise points to sample.
+    continuous_sample_points: Possible continuous noise points to sample.
+    unstructured_noise_dims: Dimensions of the unstructured noise.
+    continuous_noise_dims: Dimensions of the continuous noise.
+
+  Returns:
+    Unstructured noise, categorical noise, continuous noise numpy arrays. Each
+    should have shape [noise_samples, ?].
+  """
+  rows, cols = noise_samples, len(categorical_sample_points)
+
+  # Take random draws for non-categorical noise, making sure they are constant
+  # across columns.
+  unstructured_noise = []
+  for _ in xrange(rows):
+    cur_sample = np.random.normal(size=[1, unstructured_noise_dims])
+    unstructured_noise.extend([cur_sample] * cols)
+  unstructured_noise = np.concatenate(unstructured_noise)
+
+  continuous_noise = []
+  for _ in xrange(rows):
+    cur_sample = np.random.choice(
+        continuous_sample_points, size=[1, continuous_noise_dims])
+    continuous_noise.extend([cur_sample] * cols)
+  continuous_noise = np.concatenate(continuous_noise)
+
+  # Increase categorical noise from left to right, making sure they are constant
+  # across rows.
+  categorical_noise = np.tile(categorical_sample_points, rows)
+
+  return unstructured_noise, categorical_noise, continuous_noise
+
+
+def get_eval_noise_continuous_dim1(
+    noise_samples, categorical_sample_points, continuous_sample_points,
+    unstructured_noise_dims, continuous_noise_dims):  # pylint:disable=unused-argument
+  """Create noise showing impact of first dim continuous noise in InfoGAN.
+
+  First dimension of continuous noise is constant across columns. Other noise is
+  constant across rows.
+
+  Args:
+    noise_samples: Number of non-categorical noise samples to use.
+    categorical_sample_points: Possible categorical noise points to sample.
+    continuous_sample_points: Possible continuous noise points to sample.
+    unstructured_noise_dims: Dimensions of the unstructured noise.
+    continuous_noise_dims: Dimensions of the continuous noise.
+
+  Returns:
+    Unstructured noise, categorical noise, continuous noise numpy arrays.
+  """
+  rows, cols = noise_samples, len(continuous_sample_points)
+
+  # Take random draws for non-first-dim-continuous noise, making sure they are
+  # constant across columns.
+  unstructured_noise = []
+  for _ in xrange(rows):
+    cur_sample = np.random.normal(size=[1, unstructured_noise_dims])
+    unstructured_noise.extend([cur_sample] * cols)
+  unstructured_noise = np.concatenate(unstructured_noise)
+
+  categorical_noise = []
+  for _ in xrange(rows):
+    cur_sample = np.random.choice(categorical_sample_points)
+    categorical_noise.extend([cur_sample] * cols)
+  categorical_noise = np.array(categorical_noise)
+
+  cont_noise_dim2 = []
+  for _ in xrange(rows):
+    cur_sample = np.random.choice(continuous_sample_points, size=[1, 1])
+    cont_noise_dim2.extend([cur_sample] * cols)
+  cont_noise_dim2 = np.concatenate(cont_noise_dim2)
+
+  # Increase first dimension of continuous noise from left to right, making sure
+  # they are constant across rows.
+  cont_noise_dim1 = np.expand_dims(np.tile(continuous_sample_points, rows), 1)
+
+  continuous_noise = np.concatenate((cont_noise_dim1, cont_noise_dim2), 1)
+
+  return unstructured_noise, categorical_noise, continuous_noise
+
+
+def get_eval_noise_continuous_dim2(
+    noise_samples, categorical_sample_points, continuous_sample_points,
+    unstructured_noise_dims, continuous_noise_dims):  # pylint:disable=unused-argument
+  """Create noise showing impact of second dim of continuous noise in InfoGAN.
+
+  Second dimension of continuous noise is constant across columns. Other noise
+  is constant across rows.
+
+  Args:
+    noise_samples: Number of non-categorical noise samples to use.
+    categorical_sample_points: Possible categorical noise points to sample.
+    continuous_sample_points: Possible continuous noise points to sample.
+    unstructured_noise_dims: Dimensions of the unstructured noise.
+    continuous_noise_dims: Dimensions of the continuous noise.
+
+  Returns:
+    Unstructured noise, categorical noise, continuous noise numpy arrays.
+  """
+  rows, cols = noise_samples, len(continuous_sample_points)
+
+  # Take random draws for non-first-dim-continuous noise, making sure they are
+  # constant across columns.
+  unstructured_noise = []
+  for _ in xrange(rows):
+    cur_sample = np.random.normal(size=[1, unstructured_noise_dims])
+    unstructured_noise.extend([cur_sample] * cols)
+  unstructured_noise = np.concatenate(unstructured_noise)
+
+  categorical_noise = []
+  for _ in xrange(rows):
+    cur_sample = np.random.choice(categorical_sample_points)
+    categorical_noise.extend([cur_sample] * cols)
+  categorical_noise = np.array(categorical_noise)
+
+  cont_noise_dim1 = []
+  for _ in xrange(rows):
+    cur_sample = np.random.choice(continuous_sample_points, size=[1, 1])
+    cont_noise_dim1.extend([cur_sample] * cols)
+  cont_noise_dim1 = np.concatenate(cont_noise_dim1)
+
+  # Increase first dimension of continuous noise from left to right, making sure
+  # they are constant across rows.
+  cont_noise_dim2 = np.expand_dims(np.tile(continuous_sample_points, rows), 1)
+
+  continuous_noise = np.concatenate((cont_noise_dim1, cont_noise_dim2), 1)
+
+  return unstructured_noise, categorical_noise, continuous_noise
+
+
+def get_infogan_noise(batch_size, categorical_dim, structured_continuous_dim,
+                      total_continuous_noise_dims):
+  """Get unstructured and structured noise for InfoGAN.
+
+  Args:
+    batch_size: The number of noise vectors to generate.
+    categorical_dim: The number of categories in the categorical noise.
+    structured_continuous_dim: The number of dimensions of the uniform
+      continuous noise.
+    total_continuous_noise_dims: The number of continuous noise dimensions. This
+      number includes the structured and unstructured noise.
+
+  Returns:
+    A 2-tuple of structured and unstructured noise. First element is the
+    unstructured noise, and the second is a 2-tuple of
+    (categorical structured noise, continuous structured noise).
+  """
+  # Get unstructurd noise.
+  unstructured_noise = tf.random_normal(
+      [batch_size, total_continuous_noise_dims - structured_continuous_dim])
+
+  # Get categorical noise Tensor.
+  categorical_dist = ds.Categorical(logits=tf.zeros([categorical_dim]))
+  categorical_noise = categorical_dist.sample([batch_size])
+
+  # Get continuous noise Tensor.
+  continuous_dist = ds.Uniform(-tf.ones([structured_continuous_dim]),
+                               tf.ones([structured_continuous_dim]))
+  continuous_noise = continuous_dist.sample([batch_size])
+
+  return [unstructured_noise], [categorical_noise, continuous_noise]
+
+
+def _graph_def_from_par_or_disk(filename):
+  if filename is None:
+    return tfgan.eval.get_graph_def_from_resource(MODEL_GRAPH_DEF)
+  else:
+    return tfgan.eval.get_graph_def_from_disk(filename)

research/gan/mnist_estimator/launch_jobs.sh

+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+#!/bin/bash
+#
+# This script performs the following operations:
+# 1. Downloads the MNIST dataset.
+# 2. Trains an unconditional model on the MNIST training set using a
+#    tf.Estimator.
+# 3. Evaluates the models and writes sample images to disk.
+#
+#
+# Usage:
+# cd models/research/gan/mnist_estimator
+# ./launch_jobs.sh ${git_repo}
+set -e
+
+# Location of the git repository.
+git_repo=$1
+if [[ "$git_repo" == "" ]]; then
+  echo "'git_repo' must not be empty."
+  exit
+fi
+
+# Base name for where the evaluation images will be saved to.
+EVAL_DIR=/tmp/mnist-estimator
+
+# Where the dataset is saved to.
+DATASET_DIR=/tmp/mnist-data
+
+export PYTHONPATH=$PYTHONPATH:$git_repo:$git_repo/research:$git_repo/research/gan:$git_repo/research/slim
+
+# A helper function for printing pretty output.
+Banner () {
+  local text=$1
+  local green='\033[0;32m'
+  local nc='\033[0m'  # No color.
+  echo -e "${green}${text}${nc}"
+}
+
+# Download the dataset.
+python "${git_repo}/research/slim/download_and_convert_data.py" \
+  --dataset_name=mnist \
+  --dataset_dir=${DATASET_DIR}
+
+# Run unconditional GAN.
+NUM_STEPS=1600
+Banner "Starting training GANEstimator ${NUM_STEPS} steps..."
+python "${git_repo}/research/gan/mnist_estimator/train.py" \
+  --max_number_of_steps=${NUM_STEPS} \
+  --eval_dir=${EVAL_DIR} \
+  --dataset_dir=${DATASET_DIR} \
+  --alsologtostderr
+Banner "Finished training GANEstimator ${NUM_STEPS} steps. See ${EVAL_DIR} for output images."

research/gan/mnist_estimator/train.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Trains a GANEstimator on MNIST data."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import numpy as np
+import scipy.misc
+import tensorflow as tf
+
+from mnist import data_provider
+from mnist import networks
+
+tfgan = tf.contrib.gan
+flags = tf.flags
+
+flags.DEFINE_integer('batch_size', 32,
+                     'The number of images in each train batch.')
+
+flags.DEFINE_integer('max_number_of_steps', 20000,
+                     'The maximum number of gradient steps.')
+
+flags.DEFINE_integer(
+    'noise_dims', 64, 'Dimensions of the generator noise vector')
+
+flags.DEFINE_string('dataset_dir', None, 'Location of data.')
+
+flags.DEFINE_string('eval_dir', '/tmp/mnist-estimator/',
+                    'Directory where the results are saved to.')
+
+FLAGS = flags.FLAGS
+
+
+def _get_train_input_fn(batch_size, noise_dims, dataset_dir=None,
+                        num_threads=4):
+  def train_input_fn():
+    with tf.device('/cpu:0'):
+      images, _, _ = data_provider.provide_data(
+          'train', batch_size, dataset_dir, num_threads=num_threads)
+    noise = tf.random_normal([batch_size, noise_dims])
+    return noise, images
+  return train_input_fn
+
+
+def _get_predict_input_fn(batch_size, noise_dims):
+  def predict_input_fn():
+    noise = tf.random_normal([batch_size, noise_dims])
+    return noise
+  return predict_input_fn
+
+
+def main(_):
+  # Initialize GANEstimator with options and hyperparameters.
+  gan_estimator = tfgan.estimator.GANEstimator(
+      generator_fn=networks.unconditional_generator,
+      discriminator_fn=networks.unconditional_discriminator,
+      generator_loss_fn=tfgan.losses.wasserstein_generator_loss,
+      discriminator_loss_fn=tfgan.losses.wasserstein_discriminator_loss,
+      generator_optimizer=tf.train.AdamOptimizer(0.001, 0.5),
+      discriminator_optimizer=tf.train.AdamOptimizer(0.0001, 0.5),
+      add_summaries=tfgan.estimator.SummaryType.IMAGES)
+
+  # Train estimator.
+  train_input_fn = _get_train_input_fn(
+      FLAGS.batch_size, FLAGS.noise_dims, FLAGS.dataset_dir)
+  gan_estimator.train(train_input_fn, max_steps=FLAGS.max_number_of_steps)
+
+  # Run inference.
+  predict_input_fn = _get_predict_input_fn(36, FLAGS.noise_dims)
+  prediction_iterable = gan_estimator.predict(predict_input_fn)
+  predictions = [prediction_iterable.next() for _ in xrange(36)]
+
+  # Nicely tile.
+  image_rows = [np.concatenate(predictions[i:i+6], axis=0) for i in
+                range(0, 36, 6)]
+  tiled_image = np.concatenate(image_rows, axis=1)
+
+  # Write to disk.
+  if not tf.gfile.Exists(FLAGS.eval_dir):
+    tf.gfile.MakeDirs(FLAGS.eval_dir)
+  scipy.misc.imsave(os.path.join(FLAGS.eval_dir, 'unconditional_gan.png'),
+                    np.squeeze(tiled_image, axis=2))
+
+
+if __name__ == '__main__':
+  tf.app.run()

research/gan/mnist_estimator/train_test.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for mnist_estimator.train."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import numpy as np
+
+import tensorflow as tf
+
+import train
+
+FLAGS = tf.flags.FLAGS
+mock = tf.test.mock
+
+
+class TrainTest(tf.test.TestCase):
+
+  def test_full_flow(self):
+    FLAGS.eval_dir = self.get_temp_dir()
+    FLAGS.batch_size = 16
+    FLAGS.max_number_of_steps = 2
+    FLAGS.noise_dims = 3
+
+    # Construct mock inputs.
+    mock_imgs = np.zeros([FLAGS.batch_size, 28, 28, 1], dtype=np.float32)
+    mock_lbls = np.concatenate(
+        (np.ones([FLAGS.batch_size, 1], dtype=np.int32),
+         np.zeros([FLAGS.batch_size, 9], dtype=np.int32)), axis=1)
+    with mock.patch.object(train, 'data_provider') as mock_data_provider:
+      mock_data_provider.provide_data.return_value = (
+          mock_imgs, mock_lbls, None)
+      train.main(None)
+
+
+if __name__ == '__main__':
+  tf.test.main()