TensorFlow:在 TensorFlow 中,您可以使用 tf.config.set_visible_devices 方法来指定可见的 GPU 设备。例如,要将代码运行在 GPU 0 上,您可以在代码的开头添加以下行:
Copy
import tensorflow as tf
tf.config.set_visible_devices(tf.config.list_physical_devices('GPU')[0], 'GPU')
PyTorch:在 PyTorch 中,您可以使用 torch.cuda.set_device 方法来指定要使用的 GPU 设备。例如,要将代码运行在 GPU 0 上,您可以在代码的开头添加以下行:
Copy
import torch
torch.cuda.set_device(0)
CUDA 环境变量:另一种指定 GPU ID 的方法是使用 CUDA 环境变量。在运行代码之前,您可以设置 CUDA_VISIBLE_DEVICES 环境变量来指定要使用的 GPU ID。例如,要将代码运行在 GPU 0 上,您可以在终端中运行以下命令:
Copy
export CUDA_VISIBLE_DEVICES=0
请注意,GPU ID 从 0 开始索引,如果您有多个 GPU 设备,您可以将 GPU ID 设置为适当的值,以指定要使用的特定 GPU。