Dockerfile构建镜像

转载自：Pytorch Docker镜像构建教程（不同系统、CUDA、Python版本） - 知乎 (zhihu.com)

尽管Pytorch官方和Nvidia-NCG提供了很多版本的Pytorch镜像，但是我们平时的科研和工作中经常需要在特定系统版本、特定CUDA版本、特定Python版本下使用特定版本的Pytorch，官方的镜像很难满足自定义的需求，因此需要自己构建Pytorch镜像。下面我们来一步一步完成Pytorch镜像的构建。

 

完整的构建脚本、自动化构建流程见：cnstark/pytorch_docker: Pure Pytorch Docker Images. (github.com)
构建完成的镜像见：cnstark/pytorch (Docker Hub)

 

1 构建方法

默认读者对与Docker镜像的构建有大致的了解，我们从Dockerfile开始。

以Ubuntu20.04-Python3.9.10-CUDA11.3-Pytorch1.10.2为例，Dockerfile文件如下：

# CUDA基础镜像
FROM nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04
# CPU版本
# FROM ubuntu:20.04

# 安装基础包
RUN apt update && \
    apt install -y \
        wget build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev \
        libreadline-dev libffi-dev libsqlite3-dev libbz2-dev liblzma-dev && \
    apt clean && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /temp

# 下载python
RUN wget https://www.python.org/ftp/python/3.9.10/Python-3.9.10.tgz && \
    tar -xvf Python-3.9.10.tgz

# 编译&安装python
RUN cd Python-3.9.10 && \
    ./configure --enable-optimizations && \
    make && \
    make install

WORKDIR /workspace

RUN rm -r /temp && \
    ln -s /usr/local/bin/python3 /usr/local/bin/python && \
    ln -s /usr/local/bin/pip3 /usr/local/bin/pip

# 安装pytorch
# https://pytorch.org/get-started/locally/
RUN pip3 install torch==1.10.2+cu113 torchvision==0.11.3+cu113 torchaudio==0.10.2+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html && \
    rm -r /root/.cache/pip

 

执行

docker build -t cnstark/pytorch:1.10.2-py3.9.10-cuda11.3.1-ubuntu20.04 .

以上就完成了Pytorch Docker镜像的构建。

 

2 自动生成构建脚本

我将以上构建流程整理成了Github项目，可以自动生成构建脚本，完成构建。

2.1 克隆项目

git clone https://github.com/cnstark/pytorch_docker.git
cd pytorch_docker

2.2 生成构建脚本

执行

# CentOS8-Python3.8.10-CUDA11.1-Pytorch1.9.1
python generate_build_script.py --os centos --os-version 8 --python 3.8.10 --pytorch 1.9.1 --cuda 11.1

# Ubuntu18.04-Python3.8.10-Pytorch1.9.1 (cpu)
# python generate_build_script.py --os ubuntu --os-version 18.04 --python 3.8.10 --pytorch 1.9.1

生成如下文件

.github/workflows/docker_build_1.9.1_py3.8.10_cuda11.1_centos8.yml
scripts/build_1.9.1_py3.8.10_cuda11.1_centos8.sh

2.3 执行构建

scripts/build_1.9.1_py3.8.10_cuda11.1_centos8.sh

 

3 Github自动构建

借助Github的actions，可以在Github服务器上完成构建，并自动推送至DockerHub，避免由于本地网络等原因造成的构建失败。

上述脚本会自动生成yml文件，提交至main分支后会自动触发构建。大家如果有需要的镜像版本，可以提交issue或者把上述文件提交pull request，我merge后就会触发构建。