更改 TensorFlow 中的默认 GPU
Change default GPU in TensorFlow
根据文档,默认 GPU 是 ID 最低的那个:
If you have more than one GPU in your system, the GPU with the lowest
ID will be selected by default.
是否可以从命令行或一行代码更改此默认值?
正如 in the documentation 所述,您可以使用 tf.device('/gpu:id')
指定默认设备以外的设备。
# This will use the second GPU on your system
with tf.device('/gpu:1'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)
correctly shows how to pin your operations to a particular GPU. However, if you are running multiple TensorFlow programs on the same machine, it is recommended that you set the CUDA_VISIBLE_DEVICES
环境变量以在启动进程之前公开不同的 GPU。否则,TensorFlow 将尝试在所有可用的 GPU 上分配几乎整个内存,这会阻止其他进程使用这些 GPU(即使当前进程没有使用它们)。
注意,如果使用CUDA_VISIBLE_DEVICES
,设备名称"/gpu:0"
、"/gpu:1"
等指代0号和1号可见当前进程中的设备。
只是要清楚环境变量的使用 CUDA_VISIBLE_DEVICES
:
仅在 GPU 1 上 运行 脚本 my_script.py
,在 Linux 终端中您可以使用以下命令:
username@server:/scratch/coding/src$ CUDA_VISIBLE_DEVICES=1 python my_script.py
More 说明语法的示例:
Environment Variable Syntax Results
CUDA_VISIBLE_DEVICES=1 Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1 Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1" Same as above, quotation marks are optional
CUDA_VISIBLE_DEVICES=0,2,3 Devices 0, 2, 3 will be visible; device 1 is masked
CUDA_VISIBLE_DEVICES="" No GPU will be visible
仅供参考:
如果你想运行你的代码在第二个GPU上,假设你的机器有两个GPU,你可以做下面的技巧。
打开终端
通过键入 tmux 打开 tmux(您可以通过 sudo apt-get 安装它安装 tmux)
- 运行 tmux中的这行代码:CUDA_VISIBLE_DEVICES=1 python YourScript.py
注意:默认情况下,tensorflow 使用第一个 GPU,因此通过上述技巧,您可以 运行 在第二个 GPU 上单独编写另一个代码。
希望对您有所帮助!!
根据文档,默认 GPU 是 ID 最低的那个:
If you have more than one GPU in your system, the GPU with the lowest ID will be selected by default.
是否可以从命令行或一行代码更改此默认值?
正如 in the documentation 所述,您可以使用 tf.device('/gpu:id')
指定默认设备以外的设备。
# This will use the second GPU on your system
with tf.device('/gpu:1'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)
CUDA_VISIBLE_DEVICES
环境变量以在启动进程之前公开不同的 GPU。否则,TensorFlow 将尝试在所有可用的 GPU 上分配几乎整个内存,这会阻止其他进程使用这些 GPU(即使当前进程没有使用它们)。
注意,如果使用CUDA_VISIBLE_DEVICES
,设备名称"/gpu:0"
、"/gpu:1"
等指代0号和1号可见当前进程中的设备。
只是要清楚环境变量的使用 CUDA_VISIBLE_DEVICES
:
仅在 GPU 1 上 运行 脚本 my_script.py
,在 Linux 终端中您可以使用以下命令:
username@server:/scratch/coding/src$ CUDA_VISIBLE_DEVICES=1 python my_script.py
More 说明语法的示例:
Environment Variable Syntax Results
CUDA_VISIBLE_DEVICES=1 Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1 Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1" Same as above, quotation marks are optional
CUDA_VISIBLE_DEVICES=0,2,3 Devices 0, 2, 3 will be visible; device 1 is masked
CUDA_VISIBLE_DEVICES="" No GPU will be visible
仅供参考:
如果你想运行你的代码在第二个GPU上,假设你的机器有两个GPU,你可以做下面的技巧。
打开终端
通过键入 tmux 打开 tmux(您可以通过 sudo apt-get 安装它安装 tmux)
- 运行 tmux中的这行代码:CUDA_VISIBLE_DEVICES=1 python YourScript.py
注意:默认情况下,tensorflow 使用第一个 GPU,因此通过上述技巧,您可以 运行 在第二个 GPU 上单独编写另一个代码。
希望对您有所帮助!!