我正在尝试使用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并允许软地方,但它也没有用。
这句话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上。
缺点是您需要知道要替换的每个变量的形状。但它是可行的,至少可以让你的模型运行。