TensorFlow1.15，Estimator的内在逻辑input_fn？还是MirroredStrategy的内在逻辑？

Question

我在一台机器上用 4 个 GPU 预训练 BERT，而不是 1 个 GPU。

对于每个训练步骤，我想知道 input_fn 是给 1 个 GPU 1 个批次还是给 4 个 GPU 1 个批次。

mirrow策略代码：

    distribution = tf.contrib.distribute.MirroredStrategy(
        devices=["device:GPU:%d" % i for i in range(FLAGS.n_gpus)],
        cross_tower_ops=tf.distribute.HierarchicalCopyAllReduce())

    run_config = RunConfig(
        train_distribute=distribution,
        log_step_count_steps=log_every_n_steps,
        model_dir=FLAGS.output_dir,
        save_checkpoints_steps=FLAGS.save_checkpoints_steps)

    model_fn = model_fn_builder(
        bert_config=bert_config,
        init_checkpoint=FLAGS.init_checkpoint,
        learning_rate=FLAGS.learning_rate,
        num_train_steps=FLAGS.num_train_steps,
        num_warmup_steps=FLAGS.num_warmup_steps,
        use_tpu=FLAGS.use_tpu,
        use_one_hot_embeddings=FLAGS.use_tpu)

    estimator = Estimator(
        model_fn=model_fn,
        params={},
        config=run_config)

input_fn代码：

  def input_fn(params):
        batch_size = FLAGS.train_batch_size

        name_to_features = {
            "input_ids":
                tf.FixedLenFeature([max_seq_length], tf.int64),
            "input_mask":
                tf.FixedLenFeature([max_seq_length], tf.int64),
            "segment_ids":
                tf.FixedLenFeature([max_seq_length], tf.int64),
            "masked_lm_positions":
                tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
            "masked_lm_ids":
                tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
            "masked_lm_weights":
                tf.FixedLenFeature([max_predictions_per_seq], tf.float32),
            "next_sentence_labels":
                tf.FixedLenFeature([1], tf.int64),
        }

        if is_training:
            d = tf.data.Dataset.from_tensor_slices(tf.constant(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.contrib.data.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.repeat()

        d = d.apply(
            tf.contrib.data.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                num_parallel_batches=num_cpu_threads,
                drop_remainder=True))
        d = d.prefetch(10)
        return d

其他代码：

estimator.train(input_fn, max_steps=FLAGS.num_train_steps)

如果input_fn给1个GPU 1个batch，那么train_batch_size应该是max_batchsize_per_gpu.

如果input_fn给4个GPU 1个batch，那么train_batch_size应该是max_batchsize_per_gpu*4.

Answer 1

根据BERT-GPU。

input_fn return 1 个 GPU 的 1 个批次。

batch_size 用于 1 个 GPU。

TensorFlow1.15，Estimator的内在逻辑input_fn？还是MirroredStrategy的内在逻辑？

TensorFlow1.15, the inner logic of Estimator's input_fn? Or the inner logic of MirroredStrategy?

deep-learning

tensorflow

tensorflow-datasets

tensorflow-estimator

bert-language-model