如何在Docker构建时自动获取PyTorch与CUDA版本至环境变量

我正在编写Dockerfile，需要根据基础镜像中预装的PyTorch和CUDA版本，安装匹配的torch-scatter、torch-sparse等PyG依赖包，安装命令为：

RUN pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
RUN pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html

这里需要在构建阶段把PyTorch版本（赋值给TORCH）和CUDA版本（赋值给CUDA，注意不能带小数点）捕获为环境变量。我知道可以用ENV硬编码版本，但更换基础镜像后硬编码值会失效，导致安装错误的包。现在希望构建时自动从基础镜像中获取这两个版本并设置环境变量，该怎么做？

方法一：通过RUN命令提取并持久化环境变量 #

可以在Dockerfile中执行Python脚本提取版本，将变量写入系统环境配置文件，确保后续构建步骤和容器运行时都能读取到：

# 生成环境变量配置脚本，提取并格式化版本号
RUN echo "export TORCH=$(python -c 'import torch; print(torch.__version__)')" > /etc/profile.d/torch_env.sh && \
    echo "export CUDA=$(python -c 'import torch; print(torch.version.cuda.replace(".", ""))')" >> /etc/profile.d/torch_env.sh

# 加载变量并执行依赖安装（同一个RUN进程内完成，避免变量失效）
RUN . /etc/profile.d/torch_env.sh && \
    pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html && \
    pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html

方法二：使用ARG设置仅构建阶段可用的变量 #

如果只需要在构建过程中使用这些变量，无需保留到容器运行时，可以用ARG结合Python命令直接提取值：

# 提取PyTorch版本作为构建变量
ARG TORCH=$(python -c 'import torch; print(torch.__version__)')
# 提取CUDA版本并去掉小数点，作为构建变量
ARG CUDA=$(python -c 'import torch; print(torch.version.cuda.replace(".", ""))')

# 直接使用构建变量安装依赖
RUN pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html && \
    pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html

验证方式 #

构建镜像后，可进入容器确认版本匹配情况：

# 进入容器
docker exec -it <容器ID/名称> bash

# 检查PyTorch原始版本
python -c "import torch; print(torch.__version__)"
# 检查CUDA原始版本（带小数点）
python -c "import torch; print(torch.version.cuda)"
# 若使用方法一，可直接查看环境变量
echo $TORCH
echo $CUDA

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