获取 Tensorflow 中使用的 GPU 数量 Distributed in a multi node approach

Question

我目前正在尝试比较 Horovod 和 Tensorflow Distributed API。

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

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

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

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

我想用TF Distributed official API.

执行相同的操作并获取multi GPUs/nodes当前使用的GPU数量

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

Answer 1

回答我问题的一些发现：

等同于 hvd.size()：（与 hvd 不同，必须首先启动和初始化会话！否则你只会得到“1”） ==> tf.distribute.get_strategy().num_replicas_in_sync
等同于 hvd.rank()：（与 hvd 不同，必须首先启动和初始化会话！否则你只会得到“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 是分布式的吗运行？ : tf.distribute.has_strategy() => True/False （同上，否则你只会得到 False）

获取 Tensorflow 中使用的 GPU 数量 Distributed in a multi node approach

Get the number of GPUs used in Tensorflow Distributed in a multi node approach

python

tensorflow

nvidia

distributed