Linux下編譯安裝PyTorch源碼

• 前言
• Prerequisite
• 下載PyTorch源碼
• 編譯安裝PyTorch
• • 設定編譯線程數

本篇記錄在Linux下編譯安裝PyTorch源碼的流程，以及編譯過程中電腦當機的解決方式。

conda install astunparse numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses

使用nvcc --version查看cuda版本，在筆者機器下運行的結果為：

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Mar_21_19:15:46_PDT_2021
Cuda compilation tools, release 11.3, V11.3.58
Build cuda_11.3.r11.3/compiler.29745058_0

可以得知CUDA版本為11.3，到pytorch/repo查找並安裝對應版本的magma-cuda*：

conda install -c pytorch magma-cuda113

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
# if you are updating an existing checkout
git submodule sync
git submodule update --init --recursive --jobs 0

設定環境變量：

export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}

pytorch/From Source中給出的編譯（安裝）指令為：

python setup.py install

如果想要自行修改PyTorch中的Python源碼，可以改用以下指令。如此一來，每次改動.py檔後，不必重新安裝，改動的部分便會自動生效。

python setup.py develop

在編譯開始時，筆者的電腦對鍵盤和滑鼠的反應漸漸變得遲緩，後來又變成完全沒有反應，最後只得強制重開機解決。

如果機器的性能不是特別好，可以手動將編譯源碼的線程數調低，避免出現上述情況。

首先設定MAX_JOBS環境變量：

export MAX_JOBS=2

其中2是筆者自行選擇的數字，根據電腦的性能，也可以選擇大一點的數字。

在預期中，做完這個設定後，編譯的線程數就不會多於2了。但是ninja在編譯PyTorch中的nccl時，仍然是使用默認的6個線程，參考Compilation hangs: too many threads和Compiling NCCL does not obey MAX_JOBS, makes machine unresponsive這兩篇討論，這是由於pytorch/cmake/External/nccl.cmake中的bug造成。解決方式如下：

將pytorch/cmake/External/nccl.cmake中36,37行附近的代碼由：

        "VERBOSE=0"
        "-j"
        $ENV{MAX_JOBS}
        BUILD_BYPRODUCTS "${__NCCL_BUILD_DIR}/lib/libnccl_static.a"

修改為：

        "VERBOSE=0"
        "--jobs=2"
        BUILD_BYPRODUCTS "${__NCCL_BUILD_DIR}/lib/libnccl_static.a"

之後就能順利地編譯了。

注：其中2是筆者自行選擇的數字

筆者電腦的CPU為Intel® Core™ i5-6300HQ，GPU為GeForce GTX 950M，編譯PyTorch的過程約花了十個小時。