Ray: 如何在一台GPU上运行多个actor？

5 浏览
0 Comments

Ray: 如何在一台GPU上运行多个actor？

我只有一个GPU，想要在这个GPU上运行多个actor。以下是我使用ray库的步骤，参考链接：https://ray.readthedocs.io/en/latest/actors.html

1. 首先在GPU上定义网络

class Network():
    def __init__(self, ***一些参数***):
        self._graph = tf.Graph()
        os.environ['CUDA_VISIBLE_DIVICES'] = ','.join([str(i) for i in ray.get_gpu_ids()])
        with self._graph.as_default():
            with tf.device('/gpu:0'):
                # 在这里定义网络、损失函数和优化器
        sess_config = tf.ConfigProto(allow_soft_placement=True)
        sess_config.gpu_options.allow_growth=True
        self.sess = tf.Session(graph=self._graph, config=sess_config)
        self.sess.run(tf.global_variables_initializer())
        atexit.register(self.sess.close)
        self.variables = ray.experimental.TensorFlowVariables(self.loss, self.sess)

2. 然后定义worker类

@ray.remote(num_gpus=1)
class Worker(Network):
    # 做一些操作

3. 定义learner类

@ray.remote(num_gpus=1)
class Learner(Network):
    # 做一些操作

4. 训练函数

def train():
    ray.init(num_gpus=1)
    learner = Learner.remote(...)
    workers = [Worker.remote(...) for i in range(10)]
    # 做一些操作

当我不尝试在GPU上运行时，这个过程可以正常工作。也就是说，当我移除所有的`with tf.device('/gpu:0')`和`(num_gpus=1)`时，它可以正常工作。问题出现在我保留它们时：似乎只有`learner`被创建，而没有`workers`被构建。我该怎么做才能让它正常工作？

0