雷:如何在一个GPU上运行多个角色?

时间:2019-02-02 06:35:55

标签: python tensorflow ray

我只有一个GPU,我想在那个GPU上运行许多演员。下面是我使用的是什么ray,{以下{3}}

  1. 首先在gpu上定义网络
class Network():
    def __init__(self, ***some args here***):
        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'):
                # network, loss, and optimizer are defined here

        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)
  1. 然后定义工作者类
@ray.remote(num_gpus=1)
class Worker(Network):
    # do something
  1. 定义学习者类
@ray.remote(num_gpus=1)
class Learner(Network):
    # do something
  1. 列车功能
def train():
    ray.init(num_gpus=1)
    leaner = Learner.remote(...)
    workers = [Worker.remote(...) for i in range(10)]
    # do something

这个过程正常工作时我不试着让它在GPU上运行。也就是说,它工作正常,当我删除所有with tf.device('/gpu:0')(num_gpus=1)。当我保留它们时会出现麻烦:似乎仅创建learner,但没有构造workers。我应该怎么做才能使其正常工作?

1 个答案:

答案 0 :(得分:2)

使用装饰器@ray.remote(num_gpus=1)定义actor类时,是说从该类创建的任何actor在actor的生存期内必须为其保留一个GPU。由于只有一个GPU,因此您只能创建一个这样的actor。

如果要让多个角色共享一个GPU,则需要指定每个角色需要少于1个GPU,例如,如果要在4个角色之间共享一个GPU,则可以让每个角色都需要GPU的1/4。这可以通过使用以下方式声明actor类来实现:

@ray.remote(num_gpus=0.25)

此外,您需要确保每个角色实际上都遵守您对其施加的限制。例如,如果您想用@ray.remote(num_gpus=0.25)声明一个actor,那么您还应该确保TensorFlow最多使用GPU内存的四分之一。例如,请参见How to prevent tensorflow from allocating the totality of a GPU memory?的答案。