编译PyTorch

  1. 获取PyTorch源码。 

    1
    		 
    git clone -b v2.4.0 https://github.com/pytorch/pytorch.git --depth 1

    若出现拉取错误，可执行以下命令或重新尝试多次拉取：

    1
    		 
    git config --global http.sslBackend openssl

  2. 进入PyTorch源码。 

    1
    		 
    cd pytorch

  3. 获取子仓依赖代码。 

    1
2
    		 
    git submodule sync
git submodule update --init --recursive

  4. 设置环境变量。 

    1
2
3
4
    		 
    export USE_CUDA=1 
export USE_ROCM=1 
export USE_FBGEMM=0 
export USE_GLOO=1

  5. 安装Python依赖。 

    1
    		 
    pip3 install -r requirements.txt

  6. 编译安装PyTorch。 

    1
    		 
    python3 setup.py bdist_wheel

    安装完成后将在dist目录下生成torch*.whl文件。

  7. 安装验证。 

    1
2
    		 
    pip3 install dist/torch*.whl
pip list | grep torch

    返回信息如下：

  8. 验证是否支持GPU。 

    1
    		 
    python3

    输入如下命令：

    1
2
    		 
    import torch
print(torch.cuda.is_available())

    如果输出为“True”，则表示支持GPU。

父主题： PyTorch+TorchVision+TorchAudio源码编译指导