以下是如何在不同的环境和库中设置显卡的方法：

一、使用 CUDA 环境变量

这是最常见的方法之一，它使用CUDA环境变量来指定要使用的GPU。在运行脚本之前，设置CUDA_VISIBLE_DEVICES环境变量，将其设置为所需GPU的号码。例如：

export CUDA_VISIBLE_DEVICES=0,1,2

这将仅使设备0、1和2上的GPU可见。在Python脚本中，您可以使用以下代码来检查可见的GPU号码：

或者脚本中直接指定

import os  
#设置可见的显卡号码环境变量  
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"

或者运行脚本时设定

CUDA_VISIBLE_DEVICES=gpu_ids python3 train.py

输出可见的GPU数量

import torch  
print(torch.cuda.device_count()) 

二、使用 PyTorch 的 torch.cuda.set_device() 函数(不建议使用)

如果您使用PyTorch作为深度学习框架，可以使用torch.cuda.set_device()函数来设置当前要使用的GPU。例如：

import torch  
torch.cuda.set_device(0)  # 设置当前GPU为0号

在脚本中的其他位置，您可以使用torch.cuda.current_device()函数来检查当前选择的GPU号码：

import torch  
print(torch.cuda.current_device())  # 输出当前选择的GPU号码

三、使用 TensorFlow 的 tf.config.experimental.set_visible_devices() 函数(不建议使用)

如果您使用TensorFlow作为深度学习框架，可以使用tf.config.experimental.set_visible_devices()函数来设置可见的GPU。例如：

import tensorflow as tf  
tf.config.experimental.set_visible_devices(gpus=[0, 1, 2], per_process_gpu_memory_fraction=0.95)  # 设置可见的GPU为0、1和2号，并限制每个进程的GPU内存使用量

请注意，TensorFlow中的设置方法可能因版本而异。建议查阅TensorFlow文档以获取最新信息。

四、干货获取指南

一个以AI算法与工程技术为核心的公众号，欢迎关注交流~

以上，水平有限，有问题随时联系~