首先查看当前节点有那些可用的GPU:使用 nvidia-smi
命令来查询 Linux 服务器上可用的 GPU
nvidia-smi
命令会输出一些关于服务器上 NVIDIA GPU 的信息,包括每个 GPU 的型号、驱动版本、总内存、使用内存、温度、功率和运行在其上的进程等。
然后选择第0个和第1个GPU进行使用:
CUDA_VISIBLE_DEVICES=0,1 python train.py
设置 CUDA_VISIBLE_DEVICES
环境变量的值为 0,1
,然后运行 python train.py
命令。这意味着,当您运行训练脚本时,它只会使用服务器上的第 0 和第 1 个 GPU(假设它们是可用的)。