Python tf.distribute.OneDeviceStrategy.distribute_datasets_from_function用法及代码示例

用法

distribute_datasets_from_function(
    dataset_fn, options=None
)

分发由对 dataset_fn 的调用创建的 tf.data.Dataset 实例。

dataset_fn 将为策略中的每个工作人员调用一次。在这种情况下，我们只有一个工人和一个设备，所以 dataset_fn 被调用一次。

dataset_fn 应该采用 tf.distribute.InputContext 实例，其中可以访问有关批处理和输入复制的信息：

def dataset_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d = tf.data.Dataset.from_tensors([[1.]]).repeat().batch(batch_size)
  return d.shard(
      input_context.num_input_pipelines, input_context.input_pipeline_id)

inputs = strategy.distribute_datasets_from_function(dataset_fn)

for batch in inputs:
  replica_results = strategy.run(replica_fn, args=(batch,))

重要的：dataset_fn 返回的 tf.data.Dataset 应该具有 per-replica 批量大小，这与使用全局批量大小的 experimental_distribute_dataset 不同。这可以使用 input_context.get_per_replica_batch_size 来计算。