python - 以多节点方式获取分布式 Tensorflow 中使用的 GPU 数量

Question

我目前正在尝试比较 Horovod 和 Tensorflow 分布式 API。

使用 Horovod 时，我可以访问当前使用的 GPU 总数，如下所示:

import horovod.tensorflow as hvd
size = hvd.size()

使用 PyTorch 分布式 API 时可以使用类似的概念:

size = int(os.environ["WORLD_SIZE"])

<小时/>

我想执行相同的操作，并通过 TF Distributed 官方 API 获取多 GPU/节点当前使用的 GPU 数量。

我无法使用 CUDA_VISIBLE_DEVICES 环境变量，因为它只能在单个节点上工作。

Answer 1

一些回答我的问题的发现:

• 与 hvd.size() 等效:(与 hvd 不同，必须首先启动并初始化 session !否则您只会得到“1”)==> tf.distribute.get_strategy().num_replicas_in_sync

• 与 hvd.rank() 等效:(与 hvd 不同，必须首先启动并初始化 session !否则您只会得到“0”)

    def get_rank():
        replica_id = tf.distribute.get_replica_context().replica_id_in_sync_group
        if isinstance(replica_id, tf.Tensor):
            return tf.get_static_value(replica_id) != 0)
        else:
            return 0

• TF Distributed 是否正在运行？ : tf.distribute.has_strategy() => True/False (与上面的注释相同，否则你只会得到 False)