在TF slim中使用多个GPU用于inception_v3模型

问题描述 投票:1回答:1

我正在尝试使用3个GPU训练slim model

我特意告诉TF使用第二个GPU来分配模型:

with tf.device('device:GPU:1'):
    logits, end_points = inception_v3(inputs)

但是,每次运行代码时,我都会在GPU上出现OOM错误。我试图减少batch_size,以便模型适合内存,但网络被破坏了。

我拥有3 GPUS,有没有办法告诉TF在第二个GPU满时使用我的第三个GPU?我试过不告诉TF使用任何GPU并允许软地方,但它也没有用。

python tensorflow deep-learning distributed-computing multi-gpu
1个回答
1
投票

这句话with tf.device('device:GPU:1')告诉tensorflow专门使用GPU-1,所以它不会尝试使用你拥有的任何其他设备。

当模型太大时,recommended way将通过手动将图形拆分为不同的GPU来使用模型并行性。在您的情况下,复杂性是模型定义在库中,因此您不能为不同的层插入tf.device语句,除非您修补tensorflow。

但有一个解决方法

您可以在调用inception_v3构建器之前定义和放置变量。这样,inception_v3将重用这些变量,而不是改变它的位置。例:

with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
  with tf.device('device:GPU:1'):
    tf.get_variable("InceptionV3/Logits/Conv2d_1c_1x1/biases", shape=[1000])
    tf.get_variable("InceptionV3/Logits/Conv2d_1c_1x1/weights", shape=[1, 1, 2048, 1000])

  with tf.device('device:GPU:0'):
    logits, end_points = inception_v3(inputs)

在运行时,您将看到除Conv2d_1c_1x1之外的所有变量都放在GPU-0上,而Conv2d_1c_1x1层放在GPU-1上。

缺点是您需要知道要替换的每个变量的形状。但它是可行的,至少可以让你的模型运行。

© www.soinside.com 2019 - 2024. All rights reserved.