说明

文章是按照我解决问题的过程来写的，不是一个一步式的教程，所以显得有点凌乱。如果要操作请务必先看完全文，以免跟着中间过程走了同样的弯路。如果不想看前面的内容可以直接跳到 conda 打包这一节。

编译和安装 PyTorch（egg 格式）

PyTorch 官方只为每个 PyTorch 版本准备了几个可选的 CUDA 版本，如果需要对 PyTorch 使用不同的 CUDA 就需要自己从源码中编译。我们想要用 PyTorch 2.4.1 源码构建支持 CUDA 11.7 的包。

参考官方的说明 https://github.com/pytorch/pytorch#from-source ，过程非常简单，如果下载顺畅，在 32 核服务器上需要近一个小时编译完成。我下载的版本是 v2.4.1。总体流程是先准备好一个 conda 环境，只安装好 python 即可。然后根据说明安装各种东西，最后用 python setup.py install 或者 python setup.py develop 来安装 PyTorch。

然后 conda 环境中就会出现 PyTorch，而且用的 CUDA 版本也是系统里面的版本（如果用的镜像基于 NVIDIA 的 CUDA 镜像就肯定没问题），可以用 python -c "import torch; print(torch.version.cuda)" 来验证一下。不过，这时如果下载其他的 pip 包，pip 就会认为 torch 没有被安装，并且帮你去下载一个新的，下载完成之后 CUDA 的版本就变了！

一些想法：

1. 在构建完成 PyTorch 之后使用 pip install . 来将其安装进 lib/python3.x/site-packages 中。这个我记得是能避免 pip 去下载新的 torch 包的。
2. 在构建的时候可能需要重视对 magma-cuda* 版本的选择。（我安装的是 magma-cuda121，这并不影响我使用镜像中的 CUDA 11.7。安装之后出现了问题，不确定是否是这个导致的。现在我已经重新安装了 magma-cuda117，但好像没有什么效果？）
3. 可能得 PyTorch 构建出的版本合适。
4. 现在 egg 包已经过时了，最好转为使用 bdist_wheel。

PyTorch 编译后显示的版本不正确

一件奇怪的事：从 releases 中下载下来的 v2.4.1 版本 torch 源码用 python setup.py install 安装之后版本显示为 2.4.0a0+gitunknown。同时我的 torchvision 和 torch 也不兼容。

有个讨论 也提到了相同的问题，据说编译出来的版本都是带 git 信息的，如果需要修改版本号另行设置。不过，设置环境变量 PYTORCH_VERSION 来编译也没有成功改变版本号。实际上是环境变量名错了。

根据 另一个帖子，设置 PYTORCH_BUILD_VERSION 和 PYTORCH_BUILD_NUMBER 之后就能改变编译的版本号（两个环境变量必须同时设置），接下来又使用 python setup.py develop --cmake 来强制重新编译，用 pip install . 来安装。这样用 pip list | grep torch 就能看到版本为 2.4.1 了，再检查一下 PyTorch 配套的 CUDA 版本，和预期相符。

随后我删除了之前版本的 dist-info/（这个应该是 whl 里面的）和 egg-link 文件。

/opt/miniforge3/envs/torch241_build/lib/python3.12/site-packages/torch-2.4.0a0+gitunknown.dist-info
torch.egg-link

另外，我发现下载下来的 PyTorch 源码中有个 version.txt 文件，里面的内容为 2.4.0a0，这可能对版本也有影响。

下载的 torchvision 和当前 torch 不兼容

运行 ultralytics/yolov5/train.py 时提示：

RuntimeError: operator torchvision::nms does not exist

说明 torchvision 和 torch 不兼容的问题还是没有解决。我觉得可能是从网上下载到的 torchvision 是用的 CUDA118 或者其他版本，因此我们同样需要从源码构建 torchvision。查找对应关系后，我下载的是 v0.19.1。

先卸载 torchvision，然后按照文档安装，果然就成功了。现在 pip list | grep torch 的结果是：

torch                  2.4.1
torchvision            0.19.1a0
torchvision            0.19.1a0

为什么 torchvision 重复了？检查后发现是（用 python setup.py install 安装后） site-packages/easy-install.pth 中包含了 torchvision 的信息，同时 site-packages/ 还多了 torchvision 的 egg 文件夹，所以在 pip list 中统计了两次。

相关问题： https://github.com/pypa/setuptools/issues/3888 。构建 whl 包并安装就不会有这种 pip list 统计两次 torchvision 的情况。

有句话也不知道什么意思：

We don’t officially support building from source using pip, but if you do, you’ll need to use the --no-build-isolation flag.

现在总算是可以用了，0.19.1a0 也被认为是 0.19.1 版本，不会导致 pip 重新下载一个不同的版本。（其实这个版本号可以通过环境变量 BUILD_VERSION 来控制，见下文。）

Conda 打包

在编译了 PyTorch 之后，还有一个需求是连同着 conda 环境中其他的一些库一起打包起来。

直接打包文件夹

见 失败的尝试：导出和导入 conda 环境。

使用 conda-pack 打包报错

用 conda-pack 打包时，提示说 torch 是一个可编辑的包，不能被打包。

Collecting packages...
CondaPackError: Cannot pack an environment with editable packages
installed (e.g. from `python setup.py develop` or
 `pip install -e`). Editable packages found:

- /workspace/github/pytorch-v2.4.1

其实也容易想明白，我编译的 pytorch 产物并没有安装到 site-packages 里面。还是应该卸载掉安装的 torch，然后构建一个 whl 安装到 site-packages 里面。

卸载之前安装的开发模式的 torch

我尝试卸载 torch，但是这样安装的 torch 似乎还卸载不干净？我使用 pip uninstall torch 卸载 torch 之后 pip 仍然能识别 torch 的路径：

(torch241_build) (base) root /workspace/github/pytorch-v2.4.1 $ pip list -v | grep -v site-packages
Package                Version            Location                                                                            Installer
---------------------- ------------------ ----------------------------------------------------------------------------------- ---------
torch                  2.4.0a0+gitunknown /workspace/github/pytorch-v2.4.1

而且 torch 也确实能被正常使用。可能是 torch 并没有被完整地安装到系统里面。

先尝试 卸载已经通过 python setup.py install 安装的包：

# 在 pytorch 源码路径下记录要安装的文件
python setup.py install --record files.txt
# 删除所有相关的文件
xargs rm -rf < files.txt

结果 torch 还是存在：

(torch241_build) (base) root /workspace/github/pytorch-v2.4.1 $ pip show torch
Name: torch
Version: None
Summary: 
Home-page: 
Author: 
Author-email: 
License: 
Location: /opt/miniforge3/envs/torch241_build/lib/python3.12/site-packages
Requires: 
Required-by: thop, ultralytics, ultralytics-thop

不过 Version 变成了 None，用 pip uninstall torch 的提示也变化了，现在显示有一个 egg 可以卸载：

(torch241_build) (base) root /workspace/github/pytorch-v2.4.1 $ pip uninstall torch
Found existing installation: torch None
Uninstalling torch-None:
  Would remove:
    /opt/miniforge3/envs/torch241_build/lib/python3.12/site-packages/torch-2.4.0a0+gitunknown-py3.12.egg-info
Proceed (Y/n)? y
  Successfully uninstalled torch-None

重新卸载之后，又恢复成之前的样子，还是有个来自于 /workspace/github/pytorch-v2.4.1 的 torch 记录，torch 也还是能被正常使用。

最后到 /workspace/github/pytorch-v2.4.1 这个原来的工程里面，用 rm -rf torch.egg-info 删掉 torch.egg-info 就能从 pip 信息中成功移除 torch 了。但是这个不在标准安装路径的 torch.egg-info 是怎么被 python 找到的？而且就算是移除了，仍然可以正常导入和使用 torch 这个包。

https://stackoverflow.com/a/256614/ 还有另外一个解决方案：用 python setup.py develop -u 来解除 python setup.py develop 的行为。但是我试了没成功。

https://stackoverflow.com/a/3613880/ 指出编辑 site-packages/easy-install.pth 文件可以删除一些用开发模式安装的包：

# 删除 /workspace/github/pytorch-v2.4.1 条目
vim /opt/miniforge3/envs/torch241_build/lib/python3.12/site-packages/easy-install.pth

同样，pip show torch 是找不到了，但是 torch 还是可以在当前终端被正常使用。后来发现是因为源码目录本身就有一个名为 torch 的文件夹，所以 python 导入的 torch 并不是系统的那个。之前不在系统路径的 torch.egg-info 能被找到估计也是因为在当前路径。

换个终端 torch 就不能被使用了，不过 torch module 仍然存在可以被导入（使用时失败），这是怎么回事呢？

(torch241_build) (base) root /workspace/github/pytorch-v2.4.1 $ ls /opt/miniforge3/envs/torch241_build/lib/python3.12/site-packages/ | grep torch
functorch
torch
torchgen

原来 site-packages 里面还有文件夹残留。删了就没事了。

🚀编译并安装 torch 的 whl

设置环境变量并编译 bdist：

export PYTORCH_BUILD_VERSION=2.4.1  # 也可以是 2.4.1+cu117 这种
export PYTORCH_BUILD_NUMBER=0       # 0 和 1 好像差不多？>1 效果不一样，会有 .post 后缀

# 这个非常重要，不然编译出来的包可能不通用，不支持某些显卡
export TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.2 7.5 8.0 8.6 8.7"

# 我这里没有 NCCL，加了这个编译会失败
#export USE_SYSTEM_NCCL=1            # 默认情况下不会使用系统的 NCCL，而是会去下一个包

# 如果之前使用了错误的配置编译，就需要清理
rm -rf build                        # 好像 python setup.py clean 没什么效果

python setup.py bdist_wheel         # 注意不要写成了 bdist

完整的环境变量说明可见 https://github.com/pytorch/pytorch/blob/main/setup.py 。

安装构建好的 whl：

pip install dist/torch-2.4.0a0+gitunknown-cp312-cp312-linux_x86_64.whl

虽然日志说正在编译 2.4.1 的 bdist，但是现在文件名还是 torch-2.4.0a0+gitunknown，而且安装之后 torch 的版本也还是 2.4.0a0（后面有讲为什么）。直接修改一下元数据：

# 修改 Version 为 2.4.1
cd /opt/miniforge3/envs/torch241_build/lib/python3.12/site-packages/
vim torch-2.4.0a0+gitunknown.dist-info/METADATA
mv torch-2.4.0a0+gitunknown.dist-info/ torch-2.4.1.dist-info/

元信息修改之后版本就变了，但是最好还是把文件夹也跟着一起改了。

原来是之前搞错了，用了 python setup.py bdist（应该是 bdist_wheel），导致编译的是 dist/torch-2.4.1.linux-x86_64.tar.gz，安装的又是之前尝试时构建好的 dist/torch-2.4.0a0+gitunknown-cp312-cp312-linux_x86_64.whl，所以版本才不对。现在应该是不需要通过修改元数据来设置版本了。

🚀编译并安装 torchvision 的 whl

torchvision 的安装方式类似。不过有几点要注意：

1. 设置它的版本用的环境变量是 BUILD_VERSION，从 setup.py 可以看出来。在 CMakeLists.txt 中有个 file(STRINGS version.txt TORCHVISION_VERSION)，说明 version.txt 同样影响最终构建出来的版本。我测试发现环境变量 BUILD_VERSION 若存在，则使用它指定的版本，否则用 version.txt 中的。
2. 先安装 torch，再构建 torchvision。
3. 它也需要 TORCH_CUDA_ARCH_LIST 环境变量正确设置要编译的 CUDA 架构。不只是 torch 包需要。

安装参考：

export TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.2 7.5 8.0 8.6 8.7"
export BUILD_VERSION=0.19.1
python setup.py bdist_wheel
# Optional: install it in current environment.
pip install dist/torchvision-0.19.1-cp312-cp312-linux_x86_64.whl

🚀用 conda-pack 打包整个环境

参考 stack overflow 的帖子 和 conda 的官方帖子，如果没有 conda-pack 的话需要先安装。

conda pack -n torch241_build -o torch241_build.tar.gz

走完一段进度条就成功了。

(torch241_build) (base) root /workspace/github/pytorch-v2.4.1 $ conda pack -n torch241_build -o torch241_build.tar.gz
Collecting packages...
Packing environment at '/opt/miniforge3/envs/torch241_build' to 'torch241_build.tar.gz'
[########################################] | 100% Completed | 11min  9.7s