获取多节点方法中Tensorflow Distributed中使用的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那样！否则您将只获得“ 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）