一、PyTorch显卡禁用：场景与实现方法

在深度学习开发中，禁用显卡（GPU）的需求可能源于多种场景：

1. 调试与测试：在CPU环境下验证模型逻辑，避免GPU计算带来的不确定性；
2. 资源限制：无GPU环境或GPU被其他任务占用时，需强制使用CPU；
3. 多环境兼容：确保代码在无GPU的机器上也能运行。

PyTorch通过torch.device和CUDA_VISIBLE_DEVICES环境变量实现显卡禁用。

1.1 通过代码禁用显卡

PyTorch的核心逻辑是显式指定设备类型。若需禁用GPU，可将模型和数据强制放在CPU上：

import torch
# 禁用GPU，强制使用CPU
device = torch.device("cpu")
# 示例：定义模型并移动到CPU
model = torch.nn.Linear(10, 2).to(device)  # 模型在CPU上运行
input_data = torch.randn(5, 10).to(device)  # 输入数据也在CPU上
output = model(input_data)  # 计算在CPU上完成

此方法通过torch.device("cpu")显式指定设备，确保所有计算在CPU上进行。

1.2 通过环境变量禁用显卡

更底层的控制方式是设置CUDA_VISIBLE_DEVICES环境变量。将其设为空或无效值可完全屏蔽GPU：

# Linux/MacOS
export CUDA_VISIBLE_DEVICES=""
# Windows (PowerShell)
$env:CUDA_VISIBLE_DEVICES=""

设置后，PyTorch将无法检测到任何GPU，即使代码中未显式指定device="cpu"，也会自动回退到CPU。

1.3 验证显卡是否被禁用

通过torch.cuda.is_available()可检查CUDA是否可用：

if not torch.cuda.is_available():
    print("GPU已被禁用，当前使用CPU")
else:
    print("GPU可用")

此方法适用于调试阶段，快速确认运行环境。

二、PyTorch显卡支持：多显卡配置与优化

PyTorch对显卡的支持不仅限于单GPU，还涵盖多GPU训练、分布式计算等高级场景。

2.1 检查显卡支持情况

使用torch.cuda.device_count()可获取可用GPU数量：

num_gpus = torch.cuda.device_count()
print(f"可用GPU数量: {num_gpus}")

若输出为0，则说明无GPU可用（可能因禁用或环境问题）。

2.2 指定使用特定显卡

在多GPU环境下，可通过CUDA_VISIBLE_DEVICES或torch.cuda.set_device()指定使用的GPU：

# 方法1：通过环境变量（启动前设置）
# export CUDA_VISIBLE_DEVICES="0"  # 仅使用第0块GPU
# 方法2：通过代码指定
torch.cuda.set_device(0)  # 设置当前线程使用的GPU为第0块
# 验证当前GPU
current_device = torch.cuda.current_device()
print(f"当前使用的GPU: {current_device}")

2.3 多GPU训练配置

PyTorch提供DataParallel和DistributedDataParallel（DDP）两种多GPU训练方式。

2.3.1 DataParallel（简单多GPU）

适用于单机多卡，通过nn.DataParallel包装模型：

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)  # 包装模型
model.to("cuda")  # 移动到GPU（默认使用所有可见GPU）

注意：DataParallel存在通信开销，适合GPU数量较少（如2-4块）的场景。

2.3.2 DistributedDataParallel（高效多GPU）

适用于大规模分布式训练，需结合torch.distributed初始化：

import torch.distributed as dist
# 初始化进程组（需在代码开头调用）
dist.init_process_group(backend="nccl")  # NCCL后端适用于GPU通信
# 指定当前进程使用的GPU
local_rank = int(os.environ["LOCAL_RANK"])  # 从环境变量获取
torch.cuda.set_device(local_rank)
# 包装模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

优势：DDP通过独立进程减少通信瓶颈，适合8块以上GPU或跨节点训练。

三、常见问题与解决方案

3.1 禁用GPU后性能下降

问题：禁用GPU后训练速度显著降低。
解决方案：

• 优化CPU计算：使用num_threads参数控制PyTorch的CPU线程数（torch.set_num_threads(4)）；
• 减少批量大小（batch size）：CPU内存有限，需降低单次计算量。

3.2 多GPU训练报错

问题：DataParallel报错“CUDA error: device-side assert triggered”。
可能原因：

• GPU索引越界（如仅1块GPU但代码尝试使用第2块）；
• 模型输入数据未正确移动到GPU。
  解决方案：
• 检查torch.cuda.device_count()输出；
• 确保数据和模型在同一设备上：

  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  model.to(device)
  data = data.to(device)  # 输入数据也需移动

3.3 混合使用CPU和GPU

场景：部分操作需在CPU上完成（如数据预处理），部分在GPU上计算。
实现方法：显式移动张量：

cpu_tensor = torch.randn(10)  # 默认在CPU上
gpu_tensor = cpu_tensor.to("cuda")  # 移动到GPU
result = gpu_tensor * 2  # GPU计算
final_result = result.to("cpu")  # 移回CPU

注意：频繁的设备间数据传输会成为性能瓶颈，应尽量减少。

四、总结与最佳实践

1. 禁用GPU：优先使用torch.device("cpu")或环境变量CUDA_VISIBLE_DEVICES；
2. 多GPU训练：
   • 小规模（2-4块GPU）：DataParallel；
   • 大规模（8块+）：DistributedDataParallel；
3. 调试技巧：始终通过torch.cuda.is_available()和device检查确认运行环境；
4. 性能优化：减少CPU-GPU数据传输，合理设置批量大小和线程数。

通过掌握PyTorch的显卡禁用与支持机制，开发者可灵活应对不同硬件环境，提升代码的健壮性和可移植性。