### 如何在 Python 中指定使用特定的 GPU 进行模型训练 #### TensorFlow 指定 GPU 方法 TensorFlow 提供了多种方式来控制程序使用的 GPU 设备。通过 `tf.config` 和环境变量可以实现这一目标。 以下是具体方法： 1. **设置 CUDA_VISIBLE_DEVICES** 使用环境变量 `CUDA_VISIBLE_DEVICES` 来限定可用的 GPU 列表。例如，如果只想让程序看到第 0 号 GPU，则可以在脚本开头加入以下代码： ```python import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" ``` 2. **动态分配 GPU 资源并手动选择设备** 如果需要更灵活的方式，可以通过 TensorFlow 的 API 手动配置逻辑设备映射到物理设备上。 ```python import tensorflow as tf gpus = tf.config.list_physical_devices('GPU') if gpus: try: # 设置仅允许访问编号为 0 的 GPU tf.config.set_visible_devices(gpus[0], 'GPU') logical_gpus = tf.config.list_logical_devices('GPU') print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU") except RuntimeError as e: print(e) ``` 此外，在定义模型时也可以显式地将计算放在某个具体的 GPU 上[^1]。 #### PyTorch 指定 GPU 方法 PyTorch 同样支持多样的 GPU 控制策略，主要依赖于 `.cuda()` 函数及其变体形式。 1. **设定默认 GPU 编号** 默认情况下，PyTorch 将会优先选用第一个可见的 GPU (即索引为 0)，但如果想切换至其他 GPU，可通过如下操作完成： ```python device = torch.device("cuda:0") # 替换为 cuda:X, X 表示 GPU ID model.to(device) # 把模型移动到指定 GPU 上 input_tensor = data.cuda() # 数据也需转移到对应 GPU output = model(input_tensor) # 前向传播过程自动发生在选定 GPU 上 ``` 2. **利用 DataParallel 实现跨多个 GPU 并行处理** 当拥有不止一块 GPU 卡时，可考虑采用数据并行模式加速运算速度。不过需要注意的是，这种方式可能会带来额外开销，并不适合所有场景。 ```python from torch.nn.parallel import DataParallel model = YourModel() model = DataParallel(model).to(torch.device("cuda")) # 自动检测所有可用 GPU ``` 3. **借助 with torch.cuda.device_context(context_id)** 对某些特殊需求而言，可能希望临时改变当前活动中的 GPU，默认行为不会受到影响。 ```python with torch.cuda.device(1): # 更改为第二个 GPU second_model = AnotherModel().cuda() result = second_model(inputs) ``` 以上就是分别针对 Tensorflow 和 Pytorch 下如何指派特定 GPU 展开讨论的内容[^2][^3]。 ```python import os os.environ['CUDA_VISIBLE_DEVICES'] = '1' import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = SomeNeuralNetworkModel() model.to(device) input_data = some_input_tensor.to(device) output = model(input_data) ```