train_and_evaluate() 批量大小与 GCMLE 上的 TPU

Question

我正在尝试使用 TPUEstimator 和 train_and_evaluate() 进行 GCMLE 实验。 TPUEstimator 有一个必需的参数train_batch_size，它显然指定了批量大小。但是，对于train_and_evaluate()，我还通过 TrainSpec 指定了批量大小：

train_input = lambda: input_fn(
    filenames = hparams.train_files,
    batch_size = hparams.train_batch_size,
    hparams = hparams,
    num_epochs = hparams.num_epochs, 
    shuffle=True,
    skip_header_lines=1
    )

train_spec = tf.estimator.TrainSpec(train_input, max_steps = hparams.train_steps)

estimator = tpu_estimator.TPUEstimator(
    use_tpu=True,
    model_fn=model_fn,
    config=run_config,
    train_batch_size = hparams.train_batch_size,
    eval_batch_size = hparams.eval_batch_size,
    )
tf.estimator.train_and_evaluate(tpu_estimator, train_spec, eval_spec)

在此示例中，请考虑 train_spec 中的 train_input 指定了自己的 batch_size（例如 tf.train.batch() 或 tf.datasets.batch()）并且 train_batch_size 是 TPUEstimator 的要求.

在两个不同的地方传递train_batch_size 对我来说似乎很草率——建议只是为了确保将相同的批量大小传递给 TPUEstimator 和 TrainSpec？如果 TPUEstimator 中的 batch_size 与传递给 train_and_evaluate() 的 TrainSpec 中的 batch_size 不同，会优先考虑什么？有没有更好的方法来使用带有 TPUEstimator 的 train_and_evaluate() 并且不需要在两个不同的地方传递这个 batch_size？

此外，TPUEstimator 似乎自动创建了 params['batch_size']，根据文档，这似乎是“有效批量大小”。有效批量大小与 train_batch_size 有何关系？如果我的 train_batch_size 是 1024，是 “有效批量大小”128（因为 8 个内核）？

Answer 1

普通 Estimator 和 TPUEstimator 的批量大小处理略有不同。

对于普通的 Estimator，批量大小对 Estimator 不明确可见；相反，它是 input_fn 故事的一部分，就像您的示例所做的那样。

对于 TPU，批量大小的处理方式不同。具体来说，TPUEstimator 构造函数中的“xxx_batch_size”系列，例如训练批量大小是模型的全局批量大小。通过更改 tf.contrib.tpu.TPUConfig.per_host_input_for_training，TPUEstimator 会以不同的方式调用您的 input_fn。

这里的 params['batch_size'] 是分片批量大小，由构造函数中的 train_batch_size 计算得出。

一个具体的例子是：比如说，train_batch_size 是 64，对于 Cloud TPU，

• 如果 per_host_input_for_training 为 False，则 input_fn 将在 Cloud TPU 上调用 8 次（这称为 per-core 模式）。在这种情况下，input_fn 中的 params['batch_size'] 为 64/8=8。您的模型看到的全局总批大小为 64，即上面通过 TPUEstimator 构造函数传递的 train_batch_size。

• 如果将 per_host_input_for_training 翻转为 bool true，则 input_fn 中的 params['batch_size'] 将为 64（不是 64/8），并且 input_fn 将仅调用一次。所以，全局批大小仍然是 64。

相同的 input_fn 可以在这两种情况下工作。

对于 TPU Pod，这与 params['batch_size'] 是关于每个主机的分片批量大小相同。

总结一下：

1. 全局批量大小应通过 TPUEstimator 构造函数传递。

2. input_fn 应该从 params['batch_size'] 获取分片批量大小，并尊重它来创建数据集。

希望这会有所帮助。

Answer 2

您应该分别调用train 和evaluate 而不是train_and_evaluate。 train_and_evaluate 似乎试图以不同于 train 或 evaluate 单独设置的方式设置分布式集群。