Introduction

In TensorFlow Estimator, input_fn is meant to build and return input data, not to manage training-state logic. It does not naturally receive the current global training step as a normal parameter the way model_fn receives features and labels. If your data behavior must depend on training progress, the usual solutions are to handle that logic in model_fn, use hooks, or drive phase changes from the outer training loop rather than trying to make input_fn introspect the current step directly.

What input_fn Is Designed to Do

An Estimator input_fn typically creates and returns a tf.data.Dataset.

1import tensorflow as tf
2
3def input_fn():
4    x = tf.constant([[1.0], [2.0], [3.0], [4.0]])
5    y = tf.constant([2.0, 4.0, 6.0, 8.0])
6    ds = tf.data.Dataset.from_tensor_slices(({"x": x}, y))
7    return ds.repeat().batch(2)

Its job is data access and preprocessing. It is not the normal place for training-phase control.

model_fn Knows About Global Step

If you need training-step-aware behavior, model_fn is a more natural place because it can access the global step tensor.

1import tensorflow as tf
2
3def model_fn(features, labels, mode):
4    global_step = tf.compat.v1.train.get_or_create_global_step()
5
6    x = features["x"]
7    prediction = tf.keras.layers.Dense(1)(x)
8    loss = tf.reduce_mean(tf.square(tf.squeeze(prediction, axis=1) - labels))
9
10    train_op = tf.compat.v1.train.AdamOptimizer(0.01).minimize(
11        loss,
12        global_step=global_step
13    )
14
15    return tf.estimator.EstimatorSpec(
16        mode=mode,
17        loss=loss,
18        train_op=train_op
19    )

This is the right place for schedules such as learning-rate decay or step-dependent loss weighting.

Why Step-Dependent input_fn Is Awkward

In principle, you can create input pipelines that depend on tensors such as the global step, but that is not the usual Estimator design pattern. It makes the input pipeline harder to reason about and can complicate performance, caching, and reproducibility.

If your real goal is:

• different augmentation after a threshold
• different sampling policy later in training
• curriculum learning by training phase

it is often clearer to structure training as separate phases rather than make one input_fn magically know the step.

Use Separate Training Phases

A practical pattern is to run Estimator training in stages, each with its own input_fn.

estimator.train(input_fn=input_fn_phase_one, steps=1000)
estimator.train(input_fn=input_fn_phase_two, steps=2000)

This keeps the input functions simple and makes phase changes explicit.

It also improves debuggability because each phase has a clear boundary.

Hooks Can Observe Step Progress

If you need to react to the current step during training, session hooks are often a better fit than embedding the logic inside input_fn.

1import tensorflow as tf
2
3class StepLoggingHook(tf.estimator.SessionRunHook):
4    def before_run(self, run_context):
5        step = tf.compat.v1.train.get_global_step()
6        return tf.estimator.SessionRunArgs(step)
7
8    def after_run(self, run_context, run_values):
9        print("current step:", run_values.results)

Hooks can observe progress and coordinate external behavior without making the dataset pipeline responsible for training state.

If You Truly Need Dynamic Input Behavior

For advanced use cases, you can sometimes feed external state into the input pipeline through closures or additional tensors, but this should be a deliberate exception, not the default pattern.

For example, you might construct an input function factory:

1def make_input_fn(augment_heavily):
2    def input_fn():
3        x = tf.constant([[1.0], [2.0], [3.0], [4.0]])
4        y = tf.constant([2.0, 4.0, 6.0, 8.0])
5        ds = tf.data.Dataset.from_tensor_slices(({"x": x}, y))
6        if augment_heavily:
7            ds = ds.shuffle(4)
8        return ds.repeat().batch(2)
9    return input_fn

Then switch phases from the outer training loop rather than by querying the step inside the dataset.

Common Pitfalls

The biggest mistake is trying to force input_fn to manage training-state logic that belongs in model_fn or in the orchestration around training.

Another issue is making the input pipeline depend on hidden mutable state, which makes experiments harder to reproduce.

Developers also often ask for step-aware input functions when the simpler solution is phased training with multiple explicit input_fn definitions.

Summary

• 'input_fn is primarily for building datasets, not for knowing current training step.'
• 'model_fn can access the global step and is the better place for step-dependent model logic.'
• Use hooks when you need to observe training progress.
• Use separate training phases when input behavior should change over time.
• Keep input_fn simple unless you have a strong reason to make it state-aware.