咨询：如何在Visual Studio Code中启用GPU训练YOLO模型

1. 检查CUDA及相关依赖是否正确配置 #

• 先确认你的NVIDIA显卡算力≥3.5（支持CUDA的最低要求），安装与PyTorch/TensorFlow版本匹配的CUDA Toolkit
• 安装对应版本的cuDNN，解压后将文件复制到CUDA安装目录的对应子文件夹中
• 验证安装：终端执行nvcc -V，能显示CUDA版本即为安装成功；执行python -c "import torch; print(torch.cuda.is_available())"，返回True说明PyTorch可调用GPU

2. 强制YOLO训练脚本调用GPU #

• 训练命令中指定设备参数：以YOLOv8为例，添加--device 0（0为GPU索引，多GPU可写0,1），完整命令示例：yolo train data=coco128.yaml model=yolov8n.pt epochs=10 device=0
• 自定义脚本需手动指定设备并迁移模型/数据：

  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model = model.to(device)
  for images, labels in dataloader:
      images = images.to(device)
      labels = labels.to(device)
      # 后续训练逻辑
  

3. 确认虚拟环境的GPU支持 #

• 若使用虚拟环境，检查其中的PyTorch/TensorFlow是否为GPU版本（CPU版无法调用GPU）
• 重新安装GPU版依赖：比如PyTorch可执行对应CUDA版本的安装命令，示例：pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

4. 排查GPU占用冲突 #

• 关闭其他占用GPU的程序（游戏、其他训练任务等），执行nvidia-smi查看GPU使用情况，确认无其他进程占用
• 笔记本用户需检查是否开启独显直连，部分双显卡机型需在BIOS或显卡控制面板中设置独显优先，避免核显接管计算

5. 检查YOLO版本与配置文件 #

• 升级YOLO到最新稳定版，老旧版本可能存在GPU兼容性问题
• 查看数据配置文件（如data.yaml），确保没有强制指定CPU设备的配置项

内容的提问来源于stack exchange，提问作者Terrance Cheng