evanzd/torch-openmpi

镜像概述

evanzd/torch-openmpi镜像基于pytorch/pytorch:1.6.0-cuda10.1-cudnn7-devel构建，集成OpenSSH（用于容器间通信）和OpenMPI 4.1.0（分布式计算框架），为深度学习模型的并行训练提供稳定的运行环境。

核心功能

1. 提供PyTorch 1.6.0深度学习环境，支持CUDA10.1和CuDNN7硬件加速
2. 预装OpenSSH客户端/服务器，配置免StrictHostKeyChecking以简化容器间SSH连接
3. 集成OpenMPI 4.1.0，支持多容器/多节点的分布式计算任务

使用场景

• 深度学习模型的分布式并行训练（如大型图像分类、自然语言处理模型）
• 需要MPI通信的多容器协同计算任务
• 基于PyTorch和OpenMPI的学术研究或工程实践验证

配置说明

• SSH配置：已修改/etc/ssh/ssh_config，禁用StrictHostKeyChecking，避免容器间连接时的手动确认
• OpenMPI：已完成编译安装并配置系统路径，支持mpirun/orterun等命令直接调用

部署示例

单容器启动

bash
docker run -it --rm evanzd/torch-openmpi bash

多容器MPI任务示例（两节点）

1. 启动第一个节点容器：

bash
docker run -d --name mpi-node1 evanzd/torch-openmpi sleep infinity
NODE1_IP=$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' mpi-node1)

2. 在第二个容器中运行MPI分布式任务：

bash
docker run -it --rm evanzd/torch-openmpi mpirun -n 2 --hosts $NODE1_IP,localhost python -c "from mpi4py import MPI; comm=MPI.COMM_WORLD; print(f'Rank {comm.rank}/{comm.size}')"

（注：需提前在镜像或容器内安装mpi4py：pip install mpi4py）