elliottzheng/pytorch-mpi

PyTorch 1.10.0 开发镜像（CUDA 11.1/cuDNN 8/OpenMPI）

镜像概述

pytorch/pytorch:1.10.0-cuda11.1-cudnn8-devel+openmpi 是PyTorch官方发布的开发环境镜像，基于PyTorch 1.10.0版本构建，集成了CUDA 11.1、cuDNN 8加速库及OpenMPI通信框架。该镜像为深度学习开发提供了完整的运行时与开发工具链，支持GPU加速计算和分布式训练，适用于模型开发、调试、训练及研究场景。

核心功能与特性

基础环境

• PyTorch版本：1.10.0（稳定版，包含核心深度学习功能）
• CUDA支持：CUDA 11.1（支持NVIDIA GPU硬件加速）
• cuDNN优化：cuDNN 8（NVIDIA深度神经网络加速库，提升模型训练/推理性能）

开发工具

• Devel版本特性：包含完整的开发工具链（如编译器、头文件、库文件），支持PyTorch扩展模块编译
• 系统依赖：预安装常用系统库（如libc6-dev、git、wget等），减少环境配置工作量

分布式训练支持

• OpenMPI集成：内置OpenMPI（消息传递接口），支持多节点/多GPU分布式训练
• PyTorch分布式组件：兼容torch.distributed模块，可直接用于分布式训练脚本开发

使用场景与适用范围

• 深度学习模型开发：提供完整的PyTorch开发环境，支持模型代码编写、调试和测试
• GPU加速训练：利用CUDA 11.1和cuDNN 8优化，实现GPU加速的模型训练
• 分布式训练：通过OpenMPI支持多节点或单节点多GPU的分布式训练场景
• 学术研究：适用于需要复现PyTorch 1.10.0版本相关研究成果的场景
• 企业级应用：可作为深度学习项目的基础开发环境，确保环境一致性

使用方法与配置说明

镜像拉取

bash
docker pull pytorch/pytorch:1.10.0-cuda11.1-cudnn8-devel+openmpi

基本运行（交互式开发）

bash
docker run -it --gpus all \
  --name pytorch-dev \
  pytorch/pytorch:1.10.0-cuda11.1-cudnn8-devel+openmpi \
  /bin/bash

• --gpus all：启用所有GPU设备（需Docker支持GPU，且安装nvidia-docker）
• -it：交互式终端模式，支持命令行操作
• --name：指定容器名称，便于管理

挂载数据卷（本地代码/数据映射）

bash
docker run -it --gpus all \
  -v /path/to/local/code:/workspace/code \
  -v /path/to/local/data:/workspace/data \
  pytorch/pytorch:1.10.0-cuda11.1-cudnn8-devel+openmpi \
  /bin/bash

• -v：将本地目录挂载到容器内，实现代码和数据共享

分布式训练示例（单节点多GPU）

假设训练脚本为train.py，使用2个GPU进行分布式训练：

bash
docker run --gpus all \
  -v /path/to/train.py:/workspace/train.py \
  pytorch/pytorch:1.10.0-cuda11.1-cudnn8-devel+openmpi \
  mpirun -n 2 python /workspace/train.py

• mpirun -n 2：启动2个进程（对应2个GPU）
• 需确保训练脚本中使用torch.distributed初始化分布式环境

环境变量配置

常用环境变量可通过-e参数设置：

bash
docker run -it --gpus all \
  -e PYTHONPATH=/workspace/code \
  -e CUDA_VISIBLE_DEVICES=0,1 \  # 指定使用的GPU设备
  pytorch/pytorch:1.10.0-cuda11.1-cudnn8-devel+openmpi \
  /bin/bash

• PYTHONPATH：设置Python模块搜索路径
• CUDA_VISIBLE_DEVICES：限制容器可见的GPU设备（如仅使用第0和1号GPU）

注意事项

• GPU支持：需确保主机安装NVIDIA驱动（版本需兼容CUDA 11.1）及nvidia-docker运行时
• 版本兼容性：PyTorch 1.10.0与CUDA 11.1匹配，若需使用其他CUDA版本，需选择对应标签的镜像
• 性能优化：对于分布式训练，建议根据GPU数量调整mpirun -n参数，并优化训练脚本中的分布式通信配置