安装带CUDA10.2的PyTorch后Docker层过大，如何减小PyTorch包体积？

当前安装出现体积过大的核心原因是默认pip拉取的PyTorch包包含多CUDA版本兼容组件、冗余开发文件和调试符号，可通过以下方案逐步压缩体积：

方案1：指定CUDA版本安装专属预编译包 #

默认pip安装的通用版torch包兼容多CUDA版本，包含大量冗余组件，替换安装命令为指定CUDA 10.2的专属包，可直接减少300~500MB体积：

RUN python3 -m pip install --no-cache-dir torch==1.9.1+cu102 --extra-index-url https://download.pytorch.org/whl/cu102

方案2：安装后清理冗余文件 #

PyTorch安装目录内包含大量运行时不需要的文件，可在安装后直接删除，进一步压缩体积：

RUN python3 -m pip install --no-cache-dir torch==1.9.1+cu102 --extra-index-url https://download.pytorch.org/whl/cu102 \
    # 删除静态库文件，运行时不需要
    && find /usr/local/lib/python3.*/site-packages/torch -name "*.a" -delete \
    # 删除测试相关文件
    && find /usr/local/lib/python3.*/site-packages/torch -name "*_test.*" -delete \
    # 删除文档、示例等资源
    && rm -rf /usr/local/lib/python3.*/site-packages/torch/share \
    # 删除头文件，不需要编译自定义算子的话可直接移除
    && rm -rf /usr/local/lib/python3.*/site-packages/torch/include \
    # 剥离动态库的调试符号
    && find /usr/local/lib/python3.*/site-packages/torch/lib -name "*.so*" -exec strip {} \;

如果明确部署环境的GPU架构，还可以删除torch动态库中其他架构的CUDA编译产物，比如仅保留T4显卡对应的sm_75架构，可再减少200MB左右体积。

方案3：更换更小的基础镜像 #

不要使用带完整编译工具链的CUDA devel基础镜像，改用仅包含运行时依赖的CUDA runtime镜像，基础镜像体积可直接减少1GB以上：

# 替换前的大体积基础镜像（不推荐）
# FROM nvidia/cuda:10.2-devel-ubuntu18.04
# 替换后的精简基础镜像（推荐）
FROM nvidia/cuda:10.2-runtime-ubuntu18.04

如果不需要Ubuntu环境，也可选用Alpine体系的精简CUDA镜像，体积会更小。

方案4：使用Docker多阶段构建 #

将安装、清理PyTorch的步骤放在独立的构建阶段，最终镜像仅复制运行需要的PyTorch文件，避免构建阶段的缓存、工具链残留到最终镜像：

# 构建阶段：完成PyTorch安装、清理
FROM nvidia/cuda:10.2-devel-ubuntu18.04 as builder
RUN apt update && apt install -y python3 python3-pip
RUN python3 -m pip install --no-cache-dir torch==1.9.1+cu102 --extra-index-url https://download.pytorch.org/whl/cu102
RUN find /usr/local/lib/python3.*/site-packages/torch -name "*.a" -delete \
    && rm -rf /usr/local/lib/python3.*/site-packages/torch/share /usr/local/lib/python3.*/site-packages/torch/include \
    && find /usr/local/lib/python3.*/site-packages/torch/lib -name "*.so*" -exec strip {} \;

# 运行阶段：仅复制必要文件
FROM nvidia/cuda:10.2-runtime-ubuntu18.04
RUN apt update && apt install -y python3 && rm -rf /var/lib/apt/lists/*
COPY --from=builder /usr/local/lib/python3.*/site-packages/torch /usr/local/lib/python3.*/site-packages/torch
COPY --from=builder /usr/local/lib/python3.*/site-packages/torch-*.dist-info /usr/local/lib/python3.*/site-packages/

补充说明 #

如果部署环境不需要GPU支持，可直接安装CPU版本的PyTorch，整体体积仅100MB左右，安装命令为：

RUN python3 -m pip install --no-cache-dir torch==1.9.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu

以上方案组合使用后，PyTorch对应的Docker层体积可压缩到300~500MB区间。

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