tf.distribute.cluster_resolver.TPUClusterResolver

Резолвер кластера для TPU Google Cloud.

Наследует от: ClusterResolver

tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=None,
    zone=None,
    project=None,
    job_name='worker',
    coordinator_name=None,
    coordinator_address=None,
    credentials='default',
    service=None,
    discovery_url=None
)

Используется в блокнотах

Это реализация резолверов кластера для сервиса TPU Google Cloud.

TPUClusterResolver поддерживает следующие среды: Google Compute Engine, Google Kubernetes Engine и внутренняя среда Google.

Его можно передать в tf.distribute.TPUStrategy для поддержки обучения TF2 на TPU Cloud.

cluster_spec = tf.train.ClusterSpec({
    "ps": ["localhost:2222", "localhost:2223"],
    "worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
})

# SimpleClusterResolver is used here for illustration; other cluster
# resolvers may be used for other source of task type/id.
simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
                                        task_id=0)

...

if cluster_resolver.task_type == 'worker' and cluster_resolver.task_id == 0:
  # Perform something that's only applicable on 'worker' type, id 0. This
  # block will run on this particular instance since we've specified this
  # task to be a 'worker', id 0 in above cluster resolver.
else:
  # Perform something that's only applicable on other ids. This block will
  # not run on this particular instance.

cluster_spec = tf.train.ClusterSpec({
    "ps": ["localhost:2222", "localhost:2223"],
    "worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
})

# SimpleClusterResolver is used here for illustration; other cluster
# resolvers may be used for other source of task type/id.
simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
                                        task_id=1)

...

if cluster_resolver.task_type == 'worker':
  # Perform something that's only applicable on workers. This block
  # will run on this particular instance since we've specified this task to
  # be a worker in above cluster resolver.
elif cluster_resolver.task_type == 'ps':
  # Perform something that's only applicable on parameter servers. This
  # block will not run on this particular instance.

Методы

cluster_spec

Посмотреть исходный код

cluster_spec()

Возвращает объект ClusterSpec на основе актуальных сведений о TPU.

При каждом вызове этого метода сведения извлекаются с помощью API GCE.

connect

Посмотреть исходный код

@staticmethod
connect(
    tpu=None, zone=None, project=None
)

Инициализирует TPU и возвращает TPUClusterResolver.

Этот API подключается к удаленному кластеру TPU и инициализирует аппаратное обеспечение TPU. Пример использования:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
    tpu='')

Этот метод можно рассматривать как удобную оболочку для следующего кода:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

get_coordination_service_leader

Посмотреть исходный код

get_coordination_service_leader()

Возвращает расположение службы координации.

Служба координации должна находиться на TPU worker0.

get_job_name

Посмотреть исходный код

get_job_name()

get_master

Посмотреть исходный код

get_master()

get_tpu_system_metadata

Посмотреть исходный код

get_tpu_system_metadata()

Возвращает метаданные системы TPU.

С помощью этого метода можно получить некоторые сведения о системе TPU, например общее число ядер, число рабочих узлов TPU и устройства. Например:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tpu_system_metadata = resolver.get_tpu_system_metadata()
num_hosts = tpu_system_metadata.num_hosts

master

Посмотреть исходный код

master(
    task_type=None, task_id=None, rpc_layer=None
)

Возвращает строку Master для использования в сеансе.

В обычном случае возвращается путь grpc (grpc://1.2.3.4:8470) первого экземпляра в ClusterSpec, возвращенном функцией cluster_spec.

Если при создании TPUClusterResolver используется имя, не являющееся именем TPU, вместо него возвращается это имя (например, если аргумент tpus при создании этого TPUClusterResolver имел значение 'grpc://10.240.1.2:8470', будет возвращено 'grpc://10.240.1.2:8470').

num_accelerators

Посмотреть исходный код

num_accelerators(
    task_type=None, task_id=None, config_proto=None
)

Возвращает число ядер TPU на один рабочий узел.

Метод подключается к главному узлу, получает список всех присутствующих на нем устройств и подсчитывает их. Кроме того, перед возвратом числа ядер TPU на один узел он проверяет, что число устройств на каждом узле кластера одинаково.

set_tpu_topology

Посмотреть исходный код

set_tpu_topology(
    serialized_tpu_topology
)

Задает сведения о топологии TPU, сохраненные в этом резолвере.

__enter__

Посмотреть исходный код

__enter__()

__exit__

Посмотреть исходный код

__exit__(
    type, value, traceback
)