Python tf.distribute.DistributedIterator.get_next_as_optional用法及代碼示例

用法

get_next_as_optional()

返回包含所有副本的下一個值的 tf.experimental.Optional。

如果tf.distribute.DistributedIterator 已到達序列末尾，則返回的tf.experimental.Optional 將沒有值。

示例用法：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
global_batch_size = 2
steps_per_loop = 2
dataset = tf.data.Dataset.range(10).batch(global_batch_size)
distributed_iterator = iter(
    strategy.experimental_distribute_dataset(dataset))
def step_fn(x):
  # train the model with inputs
  return x
@tf.function
def train_fn(distributed_iterator):
  for _ in tf.range(steps_per_loop):
    optional_data = distributed_iterator.get_next_as_optional()
    if not optional_data.has_value():
      break
    per_replica_results = strategy.run(step_fn, args=(optional_data.get_value(),))
    tf.print(strategy.experimental_local_results(per_replica_results))
train_fn(distributed_iterator)
# ([0 1], [2 3])
# ([4], [])