PyTorch、TensorFlow的GPU ID命名规则及Docker GPU分配相关技术问询

1. PyTorch与TensorFlow中的GPU ID命名规则 #

• PyTorch：
  默认以从0开始的整数索引命名GPU，比如cuda:0、cuda:1对应物理机上的第一、第二块GPU。可以通过torch.cuda.device('cuda:1')指定使用某块GPU；若设置环境变量CUDA_VISIBLE_DEVICES（如export CUDA_VISIBLE_DEVICES=2,0），PyTorch内部会将可见的GPU重新映射为0、1，此时cuda:0对应物理机的第3块GPU，cuda:1对应物理机的第1块GPU。
• TensorFlow：
  默认同样使用从0开始的整数作为GPU逻辑ID，对应物理机的GPU顺序。可以通过tf.config.list_physical_devices('GPU')查看物理设备信息；CUDA_VISIBLE_DEVICES同样会影响内部ID映射，比如设置该变量为1，TensorFlow内部只会看到一块GPU，ID为0，对应物理机的第2块GPU。

2. Docker如何为容器分配GPU #

目前主流方式是借助NVIDIA Docker工具链，操作方式如下：

• 单容器分配GPU：
  • 分配所有GPU：docker run --gpus all your-image
  • 指定特定物理GPU：docker run --gpus '"device=0,2"' your-image（分配物理ID为0和2的GPU）
  • 指定GPU数量：docker run --gpus 2 your-image（随机分配2块GPU）
• Docker Compose配置：
  在docker-compose.yml的服务配置中添加GPU资源声明：

  services:
    your-service:
      image: your-image
      deploy:
        resources:
          reservations:
            devices:
              - driver: nvidia
                count: all # 或指定数量/device IDs
                capabilities: [gpu]
  

• 手动挂载方式（不推荐）：
  直接挂载NVIDIA驱动相关的设备文件和库路径，比如：

  docker run --device /dev/nvidia0 --device /dev/nvidiactl --device /dev/nvidia-uvm -v /usr/lib/nvidia-470:/usr/lib/nvidia-470 your-image
  

  这种方式需要手动匹配驱动版本，维护成本高。

3. 在配备3块GPU的机器上运行3个Docker容器时，GPU ID是否会发生变化 #

容器内部的GPU ID是相对于容器可见的GPU集合的局部索引，而非物理机的全局ID，具体情况分两种：

• 如果给每个容器单独分配1块物理GPU（比如容器1用物理0，容器2用物理1，容器3用物理2），每个容器内部只会看到1块GPU，其ID默认是0，和物理机上的全局ID不同。
• 如果某个容器分配了多块物理GPU（比如容器1分配物理0和1），容器内部会将这两块GPU映射为0和1，对应物理机的0和1，内部ID是连续的从0开始的索引。

简单来说：容器内的GPU ID是“局部索引”，不是物理机的全局ID，是否变化取决于你给容器分配的GPU范围。

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