如何在多显卡服务器上指定GPU运行TensorFlow程序

在多显卡服务器上运行深度学习模型时，合理分配GPU资源至关重要。以下是如何指定GPU来运行TensorFlow程序的实用方法：

查看GPU状态

首先，使用以下命令查看服务器上GPU的状态：

nvidia-smi

该命令将显示所有可用的GPU，并在左上角为每个GPU分配一个编号（如0、1、2、3）。这些编号将在指定GPU时用到。

方法一：在终端中指定GPU

在执行程序时，使用CUDA_VISIBLE_DEVICES变量来限制程序可见的GPU。例如：

CUDA_VISIBLE_DEVICES=1 python your_file.py

这意味着程序只会使用编号为1的GPU。以下是不同形式的使用示例：

• CUDA_VISIBLE_DEVICES=1 仅使用1号GPU
• CUDA_VISIBLE_DEVICES=0,1 使用0号和1号GPU
• CUDA_VISIBLE_DEVICES=0,2,3 使用0号、2号和3号GPU
• CUDA_VISIBLE_DEVICES="" 屏蔽所有GPU，仅使用CPU

方法二：在Python代码中指定GPU

你也可以在Python代码中指定可见的GPU。以下代码将程序限制为使用编号为0的GPU：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

通过这些方法，你可以更灵活地管理GPU资源，提高深度学习任务的执行效率。