TFGAN是谷歌开源的一个轻量级生成对抗网络(GAN)工具库，它为开发者轻松训练 GAN 提供了基础条件，提供经过完整测试的损失函数和评估指标，同时提供易于使用的范例，这些范例展示了 TFGAN 的表达能力和灵活性。这个库被包含在了TensorFlow contrib中，可以直接通过tf来进行使用，本文通过一个简单的unconditional gan模型在MNIST数据集上进行演示。

Github：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/gan

环境

• Python 3.6
• Tensorflow-gpu 1.8.0

GAN

生成式对抗网络（GAN, Generative Adversarial Networks ）是一种深度学习模型，是近年来复杂分布上无监督学习最具前景的方法之一。模型通过框架中（至少）两个模块：生成模型（Generative Model）和判别模型（Discriminative Model）的互相博弈学习产生相当好的输出。原始 GAN 理论中，并不要求 G 和 D 都是神经网络，只需要是能拟合相应生成和判别的函数即可。但实用中一般均使用深度神经网络作为 G 和 D 。一个优秀的GAN应用需要有良好的训练方法，否则可能由于神经网络模型的自由性而导致输出不理想。

GAN

在训练的过程中固定一方，更新另一方的网络权重，交替迭代，在这个过程中，双方都极力优化自己的网络，从而形成竞争对抗，直到双方达到一个动态的平衡（纳什均衡），此时生成模型 G 恢复了训练数据的分布（造出了和真实数据一模一样的样本），判别模型再也判别不出来结果，准确率为 50%，约等于乱猜。
当固定生成网络 G 的时候，对于判别网络 D 的优化，可以这样理解：输入来自于真实数据，D 优化网络结构使自己输出 1，输入来自于生成数据，D 优化网络结构使自己输出 0；当固定判别网络 D 的时候，G 优化自己的网络使自己输出尽可能和真实数据一样的样本，并且使得生成的样本经过 D 的判别之后，D 输出高概率。

上述过程可以表述为如下公式：

TFGAN

TFGAN中的训练通常包括以下步骤：

1.指定网络的输入。
2.使用GANModel设置生成器和鉴别器。
3.使用GANLoss指定损失。
4.使用GANTrainOps创建训练操作。
5.开始训练模型。

数据

首先读入MNIST数据作为输入数据x，如下所示：

import tensorflow as tf
import tensorflow.contrib.gan as tfgan
import tensorflow.contrib.layers as layers
from tensorflow.examples.tutorials.mnist import input_data


def provide_data(batch_size, num_threads=1):
    file = "MNIST"
    # range 0~1
    mnist = input_data.read_data_sets(file, one_hot=True)

    train_data = mnist.train.images.reshape(-1, 28, 28, 1) * 255
    train_labels = mnist.train.labels

    # transfer to -1~1
    train_data = (tf.to_float(train_data) - 128.0) / 128.0

    # Creates a QueueRunner for the pre-fetching operation.
    input_queue = tf.train.slice_input_producer([train_data, train_labels], shuffle=True)
    images, labels = tf.train.batch(
            input_queue,
            batch_size=batch_size,
            num_threads=num_threads,
            capacity=5 * batch_size)

    return images, labels

生成随机噪声作为输入数据z，如下所示：

    images, _ = provide_data(batch_size, num_threads=2)
    noise = tf.random_normal([batch_size, 64])

定义模型

首先我们需要定义生成器(generator)和鉴别器(discriminator)。

generator定义如下所示，将一个一维的随机噪声通过反卷积生成通道数为1的图片数据，使用tanh是为了保持生成数据的范围与输入数据一致：

def unconditional_generator(noise, weight_decay=2.5e-5, is_training=True):
    """Simple generator to produce MNIST images.

    Args:
        noise: A single Tensor representing noise.
        weight_decay: The value of the l2 weight decay.
        is_training: If `True`, batch norm uses batch statistics. If `False`, batch
            norm uses the exponential moving average collected from population 
            statistics.

    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)):
        with tf.contrib.framework.arg_scope([layers.batch_norm], is_training=is_training,
                        zero_debias_moving_mean=True):

            net = layers.fully_connected(noise, 1024)
            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

discriminator的定义如下，是一个比较简单的二分类网络，用来判断输入的数据是生成的还是真实的:

def unconditional_discriminator(img, unused_conditioning, weight_decay=2.5e-5,
                     is_training=True):
    """Discriminator network on 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.
        is_training: If `True`, batch norm uses batch statistics. If `False`, batch
            norm uses the exponential moving average collected from population 
            statistics.

    Returns:
        Logits for the probability that the image is real.
    """
    with tf.contrib.framework.arg_scope(
        [layers.conv2d, layers.fully_connected],
        activation_fn=tf.nn.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)

        with tf.contrib.framework.arg_scope([layers.batch_norm], is_training=is_training):
            net = layers.fully_connected(net, 1024, normalizer_fn=layers.batch_norm)

        return layers.linear(net, 1)

然后使用TFGAN定义一个GAN模型：

    with tf.name_scope('model'):
        # Build the generator and discriminator.
        gan_model = tfgan.gan_model(
            generator_fn=unconditional_generator,  # you define 
            discriminator_fn=unconditional_discriminator,  # you define
            real_data=images,
            generator_inputs=noise)

设置损失函数

使用TFGAN自带的损失函数配置模型，如下所示：

    with tf.name_scope('loss'):
        # Build the GAN loss.
        gan_loss = tfgan.gan_loss(
            gan_model,
            generator_loss_fn=tfgan.losses.wasserstein_generator_loss,
            discriminator_loss_fn=tfgan.losses.wasserstein_discriminator_loss,
            gradient_penalty_weight=1.0,
            add_summaries=True)

同时可以进行自定义的损失函数，损失函数接收一个gan_model，对gan_model的output进行损失计算。如下所示：

def silly_custom_generator_loss(gan_model, add_summaries=False):
    return tf.reduce_mean(gan_model.discriminator_gen_outputs)

def silly_custom_discriminator_loss(gan_model, add_summaries=False):
    return (tf.reduce_mean(gan_model.discriminator_gen_outputs) -
            tf.reduce_mean(gan_model.discriminator_real_outputs))

custom_gan_loss = tfgan.gan_loss(
    gan_model,
    generator_loss_fn=silly_custom_generator_loss,
    discriminator_loss_fn=silly_custom_discriminator_loss)

配置训练操作

接下来使用TFGAN来配置训练操作，制定模型、损失、优化器、训练率等参数。关于check_for_unused_update_ops这个参数，由于batch_norm层的原因，如果设置为True会导致更新参数检查不一致，因此需要设置为False。

    with tf.name_scope('train'):
        # Create the train ops, which calculate gradients and apply updates to weights.
        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),
            check_for_unused_update_ops=False,
            summarize_gradients=True,
            aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)

训练模型

通过TFGAN指定训练操作与保存的文件夹，就可以很容易的开始训练过程。如下所示：

    # Run the train ops in the alternating training scheme.
    tfgan.gan_train(
        train_ops,
        hooks=[tf.train.StopAtStepHook(num_steps=max_steps)],
        logdir=train_log_dir,
        save_summaries_steps=10)

可视化评估

def test(eval_dir, checkpoint_dir):
    tf.reset_default_graph()

    if not tf.gfile.Exists(eval_dir):
        tf.gfile.MakeDirs(eval_dir)

    random_inputs = tf.random_normal([100, 64])

    with tf.variable_scope('Generator'):
        images = unconditional_generator(random_inputs, is_training=False)

    reshaped_images = tfgan.eval.image_reshaper(images[:100, ...], num_cols=10)
    uint8_images = float_image_to_uint8(reshaped_images)

    image_write_ops = tf.write_file(
          '%s/%s' % (eval_dir, 'unconditional_gan.png'),
          tf.image.encode_png(uint8_images[0]))

    tf.contrib.training.evaluate_repeatedly(
            checkpoint_dir,
            eval_ops=image_write_ops,
            hooks=[tf.contrib.training.StopAfterNEvalsHook(1)],
            max_number_of_evaluations=1)

不同Epoch阶段生成的效果如下：

Epoch:400

Epoch:1000

Epoch:2000

Epoch:3000

Epoch:4000

Epoch:5000

Loss

PS:
在实验中发现生成器的学习率对生成效果有着很大的影响，最初参考官方文档使用1e-3的学习率，发现在Epoch:200左右生成了较为模糊的数字图片，进一步训练生成图片反而全是黑色背景，参考Loss变化发现整个训练过程非常不稳定。后面改用1e-4的学习率，才达到当前的生成效果。

完整的unconditional gan代码如下所示，也可以参考官方的tutorial。

import tensorflow as tf
import tensorflow.contrib.gan as tfgan
import tensorflow.contrib.layers as layers
from tensorflow.examples.tutorials.mnist import input_data


def float_image_to_uint8(image):
    """Convert float image in [-1, 1) to [0, 255] uint8.
    Note that `1` gets mapped to `0`, but `1 - epsilon` gets mapped to 255.
    Args:
        image: An image tensor. Values should be in [-1, 1).
    Returns:
        Input image cast to uint8 and with integer values in [0, 255].
    """
    image = (image * 128.0) + 128.0

    return tf.cast(image, tf.uint8)


def provide_data(batch_size, num_threads=1):
    file = "MNIST"
    # range 0~1
    mnist = input_data.read_data_sets(file, one_hot=True)

    train_data = mnist.train.images.reshape(-1, 28, 28, 1) * 255
    train_labels = mnist.train.labels

    # transfer to -1~1
    train_data = (tf.to_float(train_data) - 128.0) / 128.0

    # Creates a QueueRunner for the pre-fetching operation.
    input_queue = tf.train.slice_input_producer([train_data, train_labels], shuffle=True)
    images, labels = tf.train.batch(
            input_queue,
            batch_size=batch_size,
            num_threads=num_threads,
            capacity=5 * batch_size)

    return images, labels


def unconditional_generator(noise, weight_decay=2.5e-5, is_training=True):
    """Simple generator to produce MNIST images.

    Args:
        noise: A single Tensor representing noise.
        weight_decay: The value of the l2 weight decay.
        is_training: If `True`, batch norm uses batch statistics. If `False`, batch
            norm uses the exponential moving average collected from population 
            statistics.

    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)):
        with tf.contrib.framework.arg_scope([layers.batch_norm], is_training=is_training,
                        zero_debias_moving_mean=True):

            net = layers.fully_connected(noise, 1024)
            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_discriminator(img, unused_conditioning, weight_decay=2.5e-5,
                     is_training=True):
    """Discriminator network on 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.
        is_training: If `True`, batch norm uses batch statistics. If `False`, batch
            norm uses the exponential moving average collected from population 
            statistics.

    Returns:
        Logits for the probability that the image is real.
    """
    with tf.contrib.framework.arg_scope(
        [layers.conv2d, layers.fully_connected],
        activation_fn=tf.nn.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)

        with tf.contrib.framework.arg_scope([layers.batch_norm], is_training=is_training):
            net = layers.fully_connected(net, 1024, normalizer_fn=layers.batch_norm)

        return layers.linear(net, 1)


def train(batch_size, max_steps, gen_lr, dis_lr, train_log_dir):
    tf.reset_default_graph()

    if not tf.gfile.Exists(train_log_dir):
        tf.gfile.MakeDirs(train_log_dir)

    # Set up the input.
    images, _ = provide_data(batch_size)
    noise = tf.random_normal([batch_size, 64])

    with tf.name_scope('model'):
        # Build the generator and discriminator.
        gan_model = tfgan.gan_model(
            generator_fn=unconditional_generator,  # you define 
            discriminator_fn=unconditional_discriminator,  # you define
            real_data=images,
            generator_inputs=noise)

    with tf.name_scope('loss'):
        # Build the GAN loss.
        gan_loss = tfgan.gan_loss(
            gan_model,
            generator_loss_fn=tfgan.losses.wasserstein_generator_loss,
            discriminator_loss_fn=tfgan.losses.wasserstein_discriminator_loss,
            gradient_penalty_weight=1.0,
            add_summaries=True)

    with tf.name_scope('train'):
        # Create the train ops, which calculate gradients and apply updates to weights.
        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),
            check_for_unused_update_ops=False,
            summarize_gradients=True,
            aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)

    # Run the train ops in the alternating training scheme.
    tfgan.gan_train(
        train_ops,
        hooks=[tf.train.StopAtStepHook(num_steps=max_steps)],
        logdir=train_log_dir,
        save_summaries_steps=10)


def test(eval_dir, checkpoint_dir):
    tf.reset_default_graph()

    if not tf.gfile.Exists(eval_dir):
        tf.gfile.MakeDirs(eval_dir)

    random_inputs = tf.random_normal([100, 64])

    with tf.variable_scope('Generator'):
        images = unconditional_generator(random_inputs, is_training=False)

    reshaped_images = tfgan.eval.image_reshaper(images[:100, ...], num_cols=10)
    uint8_images = float_image_to_uint8(reshaped_images)

    image_write_ops = tf.write_file(
          '%s/%s' % (eval_dir, 'unconditional_gan.png'),
          tf.image.encode_png(uint8_images[0]))

    tf.contrib.training.evaluate_repeatedly(
            checkpoint_dir,
            eval_ops=image_write_ops,
            hooks=[tf.contrib.training.StopAfterNEvalsHook(1)],
            max_number_of_evaluations=1)


if __name__ == '__main__':
    train(16, 5000, 1e-4, 1e-4, 'logs/')
    test('eval/', 'logs/')