How to implement tensorflow Estimator with multiple models for GAN?

TensorFlow

Estimator

GAN

Machine Learning

Multiple Models

How to implement tensorflow Estimator with multiple models for GAN?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Start Practicing Learn More

Introduction

A GAN needs two models that train against each other: a generator and a discriminator. If you are forced to use TensorFlow Estimator, the usual pattern is to build both networks inside one model_fn, compute two losses, and alternate which optimizer runs on each training step.

Why GANs Feel Awkward in Estimator

Estimator was designed around a single model that exposes one training step, one loss value, and one prediction interface. A GAN breaks that shape because it has:

two networks with different objectives
separate variable groups
alternating updates during training

That does not make Estimator unusable, but it means you have to be explicit about variable scopes and train ops. If you are starting a new TensorFlow project today, a custom Keras training loop is usually simpler. If you must stay on Estimator, the structure below is the practical route.

Build Generator and Discriminator in One `model_fn`

The example below uses tf.estimator.Estimator with simple dense networks that generate vectors instead of images. The code is short enough to run, but it still shows the important parts: separate models, separate losses, and alternating updates.

python

1import numpy as np
2import tensorflow as tf
3
4tf.compat.v1.disable_eager_execution()
5
6NOISE_DIM = 8
7DATA_DIM = 4
8
9
10def make_generator(noise):
11    x = tf.keras.layers.Dense(16, activation="relu")(noise)
12    return tf.keras.layers.Dense(DATA_DIM)(x)
13
14
15def make_discriminator(inputs):
16    x = tf.keras.layers.Dense(16, activation="relu")(inputs)
17    logits = tf.keras.layers.Dense(1)(x)
18    return logits
19
20
21def model_fn(features, labels, mode, params):
22    real_data = features["real"]
23    batch_size = tf.shape(real_data)[0]
24    noise = tf.random.normal([batch_size, NOISE_DIM])
25
26    with tf.compat.v1.variable_scope("generator"):
27        fake_data = make_generator(noise)
28
29    with tf.compat.v1.variable_scope("discriminator") as scope:
30        real_logits = make_discriminator(real_data)
31        scope.reuse_variables()
32        fake_logits = make_discriminator(fake_data)
33
34    d_loss_real = tf.reduce_mean(
35        tf.nn.sigmoid_cross_entropy_with_logits(
36            labels=tf.ones_like(real_logits),
37            logits=real_logits,
38        )
39    )
40    d_loss_fake = tf.reduce_mean(
41        tf.nn.sigmoid_cross_entropy_with_logits(
42            labels=tf.zeros_like(fake_logits),
43            logits=fake_logits,
44        )
45    )
46    d_loss = d_loss_real + d_loss_fake
47
48    g_loss = tf.reduce_mean(
49        tf.nn.sigmoid_cross_entropy_with_logits(
50            labels=tf.ones_like(fake_logits),
51            logits=fake_logits,
52        )
53    )
54
55    predictions = {"generated": fake_data}
56    if mode == tf.estimator.ModeKeys.PREDICT:
57        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
58
59    g_vars = tf.compat.v1.get_collection(
60        tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope="generator"
61    )
62    d_vars = tf.compat.v1.get_collection(
63        tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope="discriminator"
64    )
65
66    total_loss = g_loss + d_loss
67
68    if mode == tf.estimator.ModeKeys.TRAIN:
69        global_step = tf.compat.v1.train.get_or_create_global_step()
70        train_generator = tf.equal(global_step % 2, 0)
71
72        g_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=0.001)
73        d_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=0.001)
74
75        g_step = g_optimizer.minimize(g_loss, var_list=g_vars)
76        d_step = d_optimizer.minimize(d_loss, var_list=d_vars)
77
78        train_op = tf.cond(train_generator, lambda: g_step, lambda: d_step)
79        with tf.control_dependencies([train_op]):
80            step_op = tf.compat.v1.assign_add(global_step, 1)
81
82        return tf.estimator.EstimatorSpec(
83            mode=mode,
84            loss=total_loss,
85            train_op=step_op,
86        )
87
88    eval_metric_ops = {
89        "generator_loss": tf.compat.v1.metrics.mean(g_loss),
90        "discriminator_loss": tf.compat.v1.metrics.mean(d_loss),
91    }
92
93    return tf.estimator.EstimatorSpec(
94        mode=mode,
95        loss=total_loss,
96        eval_metric_ops=eval_metric_ops,
97    )
98
99
100def input_fn():
101    real_samples = np.random.normal(size=(256, DATA_DIM)).astype(np.float32)
102    dataset = tf.data.Dataset.from_tensor_slices({"real": real_samples})
103    dataset = dataset.shuffle(256).repeat().batch(32)
104    return dataset
105
106
107estimator = tf.estimator.Estimator(model_fn=model_fn)
108estimator.train(input_fn=input_fn, steps=100)

The main idea is that both models live in one graph, but their trainable variables are collected separately. Training alternates between them by checking the current step.

Keep the Variable Scopes Clean

This is the part that usually breaks first. If generator and discriminator variables are not separated cleanly, an optimizer may update the wrong network or both networks at once.

The safe pattern is:

build the generator under one scope
build the discriminator under another scope
collect variables for each scope explicitly
call each optimizer with the matching var_list

Without that discipline, GAN training becomes unstable for reasons that are hard to debug.

Prediction and Evaluation

For PREDICT, you usually expose only generated samples. For EVAL, do not expect the same clean accuracy metrics you would use in supervised learning. Generator loss and discriminator loss are more realistic monitoring signals, even though they can still be noisy.

If you need visual evaluation for image GANs, save generated samples periodically outside the Estimator metrics path. That is often more informative than watching scalar losses alone.

Common Pitfalls

The most common mistake is updating both networks on every step with the same optimizer. A GAN needs separate optimizers and separate variable lists.

Another issue is reusing layers incorrectly. If the discriminator graph is rebuilt without reuse, TensorFlow creates a second set of variables and the training logic stops matching your intent.

Estimator users also often forget that GAN training is not a natural fit for the API. If the code starts becoming overly complex, that is a sign that a custom training loop may be the better engineering choice.

Finally, keep TensorFlow version compatibility in mind. Estimator-heavy code usually depends on tf.compat.v1 behavior, especially for graph mode and legacy optimizers.

Summary

A GAN in Estimator usually means one model_fn containing two sub-models.
Separate generator and discriminator variables by scope.
Compute two losses and alternate the train op between them.
Return generated samples for prediction and loss metrics for evaluation.
If you have architectural freedom, a custom training loop is often simpler than Estimator for GANs.

How to implement tensorflow Estimator with multiple models for GAN?

Master System Design with Codemia

Introduction

Why GANs Feel Awkward in Estimator

Build Generator and Discriminator in One model_fn

Keep the Variable Scopes Clean

Prediction and Evaluation

Common Pitfalls

Summary

Build Generator and Discriminator in One `model_fn`