如何将PyTorch版本和CUDA版本保存到环境变量中

Question

我正在编写一个Dockerfile，它需要安装不同的pip轮子，具体取决于基本映像中安装的PyTorch版本和CUDA版本：

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

因此，我需要在构建时将这些版本捕获到环境变量TORCH和CUDA中。

我知道我可以在Dockerfile中使用ENV命令来分配给环境变量：

ENV TORCH=1.12.0
ENV CUDA=113

(注意，与TORCH不同，CUDA不能包含任何点(，然后构建容器。现在，如果我登录到正在运行的Docker容器，我可以从命令行获得以下版本：

python -c "import torch; print(torch.__version__)"
>>> 1.12.0
python -c "import torch; print(torch.version.cuda)"
>>> 11.3

然而，我不想对Dockerfile中的版本进行硬编码，因为如果我更改基本图像，硬编码的值将是错误的，我会尝试安装错误的轮子，安装将失败。我想在构建时找到它们，并将它们分配给TORCH&CUDA。我该怎么做？

Answer 1

您可以使用普通的shell语法在RUN命令中设置环境变量，但这些设置在该命令结束时会丢失。因此，在单个RUN命令中您可以使用shell命令替换来设置环境变量并使用它，但在该命令之后，它的值将不再可用。

# all within a single RUN line
RUN TORCH=$(python -c "import torch; print(torch.__version__)"); 
CUDA=$(python -c "import torch; print(torch.version.cuda)" | sed 's/.//g'); 
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