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

Question

我正在尝试使用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并允许软地方，但它也不起作用。

Answer 1

此声明with tf.device('device:GPU:1')告诉tensorflow专门使用GPU-1，因此它不会尝试使用您拥有的任何其他设备。

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

但有一种解决方法

您可以在调用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上。

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