AWS EKS中G4实例Docker容器运行YOLOv5 GPU训练报错求助

核心问题 #

容器内PyTorch无法识别CUDA设备，导致YOLOv5训练时--device 0参数触发断言错误。

排查与解决步骤 #

• 验证容器GPU权限
  运行容器时必须指定GPU资源：

  • Docker命令需添加--gpus all（或--gpus device=0）；
  • Kubernetes部署需在Pod配置中添加GPU资源请求：

    resources:
      requests:
        nvidia.com/gpu: 1
      limits:
        nvidia.com/gpu: 1
    

  进入容器执行以下命令验证：

  python3 -c "import torch; print(torch.cuda.is_available())"
  

  若输出False，说明容器未获取GPU权限，需调整运行配置。

• 检查PyTorch与CUDA版本兼容性
  PyTorch 1.12.1官方适配CUDA 11.3/11.6版本，与当前使用的CUDA 11.4存在版本不匹配问题。重新安装适配版本的PyTorch：

  pip3 install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
  

  安装后再次执行python3 -c "import torch; print(torch.cuda.is_available())"确认是否返回True。

• 验证EKS节点NVIDIA环境配置
  确认EKS节点已部署NVIDIA容器运行时和NVIDIA Device Plugin：
  执行kubectl describe node <你的G4节点名称>，查看Capacity和Allocatable字段中是否包含nvidia.com/gpu: 1，确保GPU资源可被Pod调度。

• 检查容器内CUDA环境
  进入容器执行nvcc -V确认CUDA版本，同时执行nvidia-smi验证是否能正常显示GPU信息。若容器内无法运行nvidia-smi，说明容器运行时未正确加载NVIDIA驱动，需检查节点上的NVIDIA容器工具包配置。

• YOLOv5设备选择测试
  若以上验证均正常，先在Python交互环境测试CUDA设备识别：

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

  能正常输出设备信息后，尝试运行python3 train.py --device cuda启动训练。

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