解决办法:使用启用最少的GPU显存来运行程序或者限定显存比例
在 TensorFlow 2.x 中,设置 GPU 的方法略有不同。以下是如何在 TensorFlow 2.x 中指定使用第二个 GPU(即 GPU 1)并启用最少的 GPU 显存来运行程序的代码示例:
-
导入 TensorFlow:
import tensorflow as tf
-
设置内存自增长:
gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: try: # 设置内存自增长 for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) print('-------------已设置完GPU内存自增长--------------') # 指定使用第二块GPU tf.config.experimental.set_visible_devices(gpus[1], 'GPU') # 获取逻辑GPU个数 logical_gpus = tf.config.experimental.list_logical_devices('GPU') print('逻辑GPU个数为:', len(logical_gpus)) except RuntimeError as e: # 打印异常信息 print(e)
-
运行示例代码:
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) c = tf.matmul(a, b) print(c)
这段代码首先获取系统中所有的物理 GPU,然后设置它们的内存自增长。接着,通过 tf.config.experimental.set_visible_devices(gpus[1], 'GPU')
指定使用第二块 GPU。最后,运行一个简单的矩阵乘法操作来验证设置是否生效。
注意:在 TensorFlow 2.x 中,tf.config.experimental
模块用于进行这些配置。确保在程序的开始处设置这些配置,以确保它们在 TensorFlow 初始化之前生效。