Tensorflow Estimator API下动态样本计数的实现方案咨询

Great question! Since you're working with TensorFlow's Estimator API and need to track total processed samples (not just optimizer steps) while supporting dynamic batch sizes, here are two robust solutions that fit your requirements:

Solution 1: Reuse global_step with Dynamic Increments #

If you don't need to keep track of optimizer steps separately, repurposing global_step to count samples is a clean approach—since Estimator automatically saves and restores this variable by default. The key is to manually update global_step with the current batch size instead of letting the optimizer increment it by 1.

Here's how to implement this in your model_fn:

def model_fn(features, labels, mode, params):
    # Define your model architecture and compute loss
    model = build_your_model()
    logits = model(features['input_data'])
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    # Get the global_step variable
    global_step = tf.train.get_global_step()
    # Dynamically get current batch size from input features
    batch_size = tf.shape(features['input_data'])[0]

    if mode == tf.estimator.ModeKeys.TRAIN:
        # Prevent optimizer from auto-incrementing global_step by setting global_step=None
        optimizer = tf.train.AdamOptimizer(learning_rate=params['lr'])
        base_train_op = optimizer.minimize(loss, global_step=None)
        
        # Update global_step with batch size as a control dependency
        with tf.control_dependencies([tf.assign_add(global_step, batch_size)]):
            train_op = base_train_op

        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
    # Handle EVAL and PREDICT modes as usual
    elif mode == tf.estimator.ModeKeys.EVAL:
        eval_metrics = {'accuracy': tf.metrics.accuracy(labels=labels, predictions=tf.argmax(logits, 1))}
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metrics)
    
    else:
        predictions = {'class': tf.argmax(logits, 1)}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

Key Notes: #

• Set global_step=None in optimizer.minimize() to avoid automatic increments of 1.
• tf.shape(features['input_data'])[0] grabs the actual batch size for each step, which handles dynamic changes perfectly.
• Estimator will still save/restore global_step exactly as it does by default—no extra configuration needed.

Solution 2: Create a Custom Variable with Automatic Save/Restore #

If you want to keep global_step for tracking optimizer steps and have a separate variable for sample count, you can create a custom global variable. Estimator automatically saves all variables in the GLOBAL_VARIABLES collection, so this variable will be persisted and restored without extra work.

Here's the implementation:

def model_fn(features, labels, mode, params):
    # Model definition and loss calculation
    model = build_your_model()
    logits = model(features['input_data'])
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    # Create a non-trainable variable to track total samples
    sample_count = tf.get_variable(
        name='total_processed_samples',
        shape=[],
        dtype=tf.int64,
        initializer=tf.zeros_initializer(),
        trainable=False  # Critical: don't let optimizer update this
    )
    batch_size = tf.shape(features['input_data'])[0]

    if mode == tf.estimator.ModeKeys.TRAIN:
        # Update sample count with current batch size
        update_sample_count = tf.assign_add(sample_count, batch_size)
        
        # Let optimizer handle global_step (for step count)
        optimizer = tf.train.AdamOptimizer(learning_rate=params['lr'])
        base_train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
        
        # Ensure sample count is updated before training step runs
        with tf.control_dependencies([update_sample_count]):
            train_op = base_train_op

        # Optional: Track sample count as a training metric
        train_metrics = {'total_samples': tf.metrics.mean(sample_count)}
        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            train_op=train_op,
            training_hooks=[tf.train.LoggingTensorHook({'total_samples': sample_count}, every_n_iter=100)]
        )
    
    # Include sample count in eval metrics if needed
    elif mode == tf.estimator.ModeKeys.EVAL:
        eval_metrics = {
            'accuracy': tf.metrics.accuracy(labels=labels, predictions=tf.argmax(logits, 1)),
            'total_processed_samples': tf.metrics.mean(sample_count)
        }
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metrics)
    
    else:
        predictions = {'class': tf.argmax(logits, 1)}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

Key Notes: #

• Setting trainable=False ensures the optimizer doesn't modify this variable—we control updates manually.
• The variable is part of the default GLOBAL_VARIABLES collection, so Estimator saves/restores it automatically.
• You can log or report this variable as a metric to monitor total samples during training/evaluation.

Final Recommendation #

• Use Solution 1 if you only need to track sample count and don't care about retaining the original "step count" semantics of global_step.
• Use Solution 2 if you need to track both optimizer steps and total processed samples separately.

Both solutions fully support dynamic batch sizes and are compatible with TensorFlow's Estimator API.

内容的提问来源于stack exchange，提问作者Simon Rühle