在 TensorFlow Estimator 中，input_fn 可以知道当前的训练步骤吗？

Question

在我的模型中（使用 Tensorflow Estimator），我希望我的数据更加动态。例如。在训练期间提供不同的数据（在不同的训练步骤中，向模型提供不同的数据）。

下面的代码示例。 get_input_fn 提供 input_fn 和一个 _parse 函数来处理这些特征。 _parse 中的 _py_process_line_pair 进行精确处理。但我不确定如何将 global_step（或相关参数传递给 _py_process_line_pair）

    def _parse(self, features):
      def _py_process_line_pair(src_wds, trg_wds, cur_training_steps):
        .... (some processing depends on cur_training_steps)
        return np.array(src_ids, np.int32), np.array(trg_ids, np.int32)

    src_wds, trg_wds = features['src_wds'], features['trg_wds']
    src_ids, trg_ids = tf.py_func(
        _py_process_line_pair,
        [src_wds, trg_wds],
        [tf.int32, tf.int32])
    src_ids.set_shape(
        [self.flags.max_src_len])
    trg_ids.set_shape(
        [self.flags.max_trg_len])
    output = {
        'src_ids': src_ids,
        'trg_ids': trg_ids,
    }
    return output

  def get_input_fn(self, is_training, input_files, num_cpu_threads):

    def input_fn(params):
        batch_size = params['batch_size']
        if is_training:
            d = tf.data.Dataset.from_tensor_slices(tf.constant(tf.gfile.Glob(input_files)))
            d = d.repeat()
            d = d.shuffle(buffer_size=len(input_files))
            cycle_length = min(num_cpu_threads, len(input_files))
            d = d.apply(
                tf.data.experimental.parallel_interleave(
                    tf.data.TFRecordDataset,
                    sloppy=is_training,
                    cycle_length=cycle_length))
            d = d.shuffle(buffer_size=100)
        else:
            d = tf.data.TFRecordDataset(input_files)

        d = d.apply(
            tf.data.experimental.map_and_batch(
                lambda record:  self._parse(tf.parse_single_example(record, self.feature_set)),
                batch_size=batch_size,
                num_parallel_batches=num_cpu_threads,
                drop_remainder=is_training))
        return d
    return input_fn

Answer 1

这很简单：您只需在 _parse 函数中使用 tf.train.get_or_create_global_step() 从图中获取 global_step 张量。

这是一个工作示例

import tensorflow as tf
import numpy as np

# Synth dataset with 10 values
x = np.arange(10)

# This function replaces 'x' by the current step
def step_dependant_preprocessing(x):
    global_step = tf.train.get_or_create_global_step()
    return global_step

# Maps step_dependant_preprocessing
def input_fn():
    dataset = tf.data.Dataset.from_tensor_slices((x))
    dataset = dataset.map(step_dependant_preprocessing)
    return dataset

def model_fn(features, labels, mode, params=None):
    # Get the global step
    global_step = tf.train.get_or_create_global_step()

    # Since this example doesn't use an optimizer, we need to increment
    # the global step manually.
    increment_global_step = tf.assign_add(global_step, 1)

    # Logging hook to verify that the global step inside the input fn has 
    # the same value as the one here.
    logging_hook = tf.train.LoggingTensorHook({"true_global_step": global_step, 
                                               "input_fn_global_step": features}, 
                                              every_n_iter=1)

    return tf.estimator.EstimatorSpec(
        mode=tf.estimator.ModeKeys.TRAIN,
        loss=tf.constant(0.0), # Needed to use estimator.train()
        training_hooks=[logging_hook],
        train_op=increment_global_step)

estimator = tf.estimator.Estimator(model_fn=model_fn)

estimator.train(input_fn)

...

# Output

INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Create CheckpointSaverHook.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Saving checkpoints for 0 into /tmp/tmppuwe9hxh/model.ckpt.
INFO:tensorflow:loss = 0.0, step = 1
INFO:tensorflow:input_fn_global_step = 1, true_global_step = 1
INFO:tensorflow:input_fn_global_step = 2, true_global_step = 2 (0.007 sec)
INFO:tensorflow:input_fn_global_step = 3, true_global_step = 3 (0.002 sec)
INFO:tensorflow:input_fn_global_step = 4, true_global_step = 4 (0.001 sec)
INFO:tensorflow:input_fn_global_step = 5, true_global_step = 5 (0.001 sec)
INFO:tensorflow:input_fn_global_step = 6, true_global_step = 6 (0.001 sec)
INFO:tensorflow:input_fn_global_step = 7, true_global_step = 7 (0.001 sec)
INFO:tensorflow:input_fn_global_step = 8, true_global_step = 8 (0.001 sec)
INFO:tensorflow:input_fn_global_step = 9, true_global_step = 9 (0.001 sec)
INFO:tensorflow:input_fn_global_step = 10, true_global_step = 10 (0.001 sec)
INFO:tensorflow:Saving checkpoints for 11 into /tmp/tmppuwe9hxh/model.ckpt.
INFO:tensorflow:Loss for final step: 0.0.