tf.distribute.InputOptions  |  TensorFlow v2.16.1 (original) (raw)

tf.distribute.InputOptions

Stay organized with collections Save and categorize content based on your preferences.

Run options for experimental_distribute_dataset(s_from_function).

tf.distribute.InputOptions(
    experimental_fetch_to_device=None,
    experimental_replication_mode=tf.distribute.InputReplicationMode.PER_WORKER,
    experimental_place_dataset_on_device=False,
    experimental_per_replica_buffer_size=1
)

Used in the notebooks

Used in the guide Used in the tutorials
Migrate from TPU embedding_columns to TPUEmbedding layer TensorFlow 2 TPUEmbeddingLayer: Quick Start

This can be used to hold some strategy specific configs.

# Setup TPUStrategy
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

dataset = tf.data.Dataset.range(16)
distributed_dataset_on_host = (
    strategy.experimental_distribute_dataset(
        dataset,
        tf.distribute.InputOptions(
            experimental_replication_mode=
            experimental_replication_mode.PER_WORKER,
            experimental_place_dataset_on_device=False,
            experimental_per_replica_buffer_size=1)))
Attributes
experimental_fetch_to_device Boolean. If True, dataset elements will be prefetched to accelerator device memory. When False, dataset elements are prefetched to host device memory. Must be False when using TPUEmbedding API. experimental_fetch_to_device can only be used with experimental_replication_mode=PER_WORKER. Default behavior is same as setting it to True.
experimental_replication_mode Replication mode for the input function. Currently, the InputReplicationMode.PER_REPLICA is only supported with tf.distribute.MirroredStrategy. experimental_distribute_datasets_from_function. The default value is InputReplicationMode.PER_WORKER.
experimental_place_dataset_on_device Boolean. Default to False. When True, dataset will be placed on the device, otherwise it will remain on the host. experimental_place_dataset_on_device=True can only be used with experimental_replication_mode=PER_REPLICA
experimental_per_replica_buffer_size Integer. Default to 1. Indicates the prefetch buffer size in the replica device memory. Users can set it to 0 to completely disable prefetching behavior, or a number greater than 1 to enable larger buffer size. Note that this option is still valid with experimental_fetch_to_device=False.

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates. Some content is licensed under the numpy license.

Last updated 2024-04-26 UTC.