anibali/pytorch

PyTorch Docker镜像

Ubuntu + PyTorch + CUDA（可选）

需求

使用此镜像需先安装Docker Engine。Docker Engine的安装说明可在Docker官方网站获取。

CUDA需求

若您有兼容CUDA的NVIDIA显卡，可使用支持CUDA的PyTorch镜像以启用硬件加速。此功能仅在Ubuntu Linux系统中测试过。

首先，确保安装了适当的NVIDIA驱动。在Ubuntu上，最简单的方法是通过NVIDIA官方CUDA下载页面安装至少与您要使用的镜像版本相同的CUDA，这将确保驱动版本正确。例如，若使用cuda-10.1镜像，安装CUDA 10.1或10.2可确保驱动正确。

还需安装NVIDIA Container Toolkit以允许Docker容器访问GPU设备，可从https://github.com/NVIDIA/nvidia-docker%E8%8E%B7%E5%8F%96%E3%80%82

预构建镜像

预构建镜像可在Docker Hub的https://hub.docker.com/r/anibali/pytorch/%E4%BB%93%E5%BA%93%E8%8E%B7%E5%8F%96%E3%80%82

例如，拉取包含PyTorch 1.13.0和CUDA 11.8的镜像：

bash
$ docker pull anibali/pytorch:1.13.0-cuda11.8

使用方法

运行PyTorch脚本

可使用python3命令在容器内运行PyTorch程序。例如，若当前目录包含一个以main.py为入口的PyTorch项目，可通过以下命令运行：

sh
docker run --rm -it --init \
  --gpus=all \
  --ipc=host \
  --user="$(id -u):$(id -g)" \
  --volume="$PWD:/app" \
  anibali/pytorch python3 main.py

上述Docker命令行选项说明：

• --gpus=all：使用CUDA时必填，否则可选。将主机的显卡传递给容器。也可通过此选项精确控制暴露的显卡（详见文档：https://github.com/NVIDIA/nvidia-docker%EF%BC%89%E3%80%82
• --ipc=host：使用多进程时必填，原因见https://github.com/pytorch/pytorch#docker-image%E3%80%82
• --user="$(id -u):$(id -g)"：设置容器内用户匹配主机用户和组ID。可选，但有助于确保文件写入权限正确。
• --volume="$PWD:/app"：将当前工作目录挂载到容器内。容器内默认工作目录为/app。可选。

运行图形化应用

若在Linux主机上运行，可让容器内代码通过主机X服务器显示图形（例如使用OpenCV的imshow）。以下是一种简单但不安全的方法。更全面的GUI与Docker指南见[***]

在主机上运行：

sh
sudo xhost +local:root

稍后可通过sudo xhost -local:root撤销权限。运行容器时，需添加-e "DISPLAY"和--volume="/tmp/.X11-unix:/tmp/.X11-unix:rw"选项，以提供X11 socket和显示ID。示例：

sh
docker run --rm -it --init \
  --gpus=all \
  -e "DISPLAY" --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" \
  anibali/pytorch python3 -c "import tkinter; tkinter.Tk().mainloop()"

构建自定义镜像

添加额外依赖的推荐方式是基于本项目的PyTorch镜像创建自定义Dockerfile。

例如，如需OpenCV并使用PyTorch 1.13.0，可创建以下Dockerfile，以anibali/pytorch:1.13.0-cuda11.8-ubuntu22.04为基础镜像并安装OpenCV：

dockerfile
FROM anibali/pytorch:1.13.0-cuda11.8-ubuntu22.04

# 设置时区
ENV TZ=UTC
RUN sudo ln -snf /usr/share/zoneinfo/$TZ /etc/localtime

# 安装OpenCV所需的系统库
RUN sudo apt-get update \
 && sudo apt-get install -y libgl1-mesa-glx libgtk2.0-0 libsm6 libxext6 \
 && sudo rm -rf /var/lib/apt/lists/*

# 从PyPI安装OpenCV
RUN pip install opencv-python==4.5.1.48

开发与贡献

dockerfiles/目录下的Dockerfile由manager.py脚本根据images.yml中的配置和templates/目录下的模板自动生成。

创建新Dockerfile的示例流程：

1. （可选）若现有模板不适用，在templates/目录创建新模板文件。
2. 在images.yml中添加新条目（参考现有条目）。
3. 运行python manager.py生成Dockerfile，新目录将在dockerfiles/下创建。
4. 构建生成的Dockerfile并测试。若仅个人使用，至此即可。
5. （可选）如认为新镜像对他人有用，可提交PR，将考虑发布。