如何在Ubuntu上编译PyTorch

在Ubuntu上编译PyTorch可以从源代码构建，以便获得最佳的性能和定制选项。以下是在Ubuntu上编译PyTorch的详细步骤：

前提条件

1. 安装依赖项：

   • 打开终端并运行以下命令来更新包列表并安装必要的依赖项：

     sudo apt update
     sudo apt install -y build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev
     

2. 安装Python和虚拟环境：

   • 安装Python 3和pip：

     sudo apt install python3 python3-pip
     

   • 创建一个虚拟环境（可选但推荐）：

     python3 -m venv pytorch-env
     source pytorch-env/bin/activate
     

3. 安装CUDA和cuDNN（如果需要GPU支持）：

   • 下载并安装CUDA Toolkit：CUDA Toolkit
   • 下载并安装cuDNN：cuDNN

编译PyTorch

1. 克隆PyTorch仓库：

   git clone --recursive https://github.com/pytorch/pytorch
   cd pytorch
   

2. 设置环境变量（如果使用GPU）：

   • 设置CUDA和cuDNN的路径：

     export CUDA_HOME=/usr/local/cuda
     export CUDNN_INCLUDE_DIR=/usr/local/cuda/include
     export CUDNN_LIB_DIR=/usr/local/cuda/lib64
     

3. 安装Python依赖项：

   pip install -r requirements.txt
   

4. 配置CMake：

   • 创建一个构建目录并进入该目录：

     mkdir build
     cd build
     

   • 运行CMake配置脚本：

     cmake .. -DCMAKE_BUILD_TYPE=Release \
           -DPYTHON_EXECUTABLE=$(which python3) \
           -DPYTHON_INCLUDE_DIR=$(python3 -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())") \
           -DPYTHON_LIBRARY=$(python3 -c "import distutils.sysconfig as sysconfig; print(sysconfig.get_config_var('LIBDIR'))")/libpython3.x.so \
           -DUSE_CUDA=ON \
           -DCUDA_ARCH_NAME=Your_GPU_Architecture \
           -DCUDA_HOME=$CUDA_HOME \
           -DCUDNN_INCLUDE_DIR=$CUDNN_INCLUDE_DIR \
           -DCUDNN_LIB_DIR=$CUDNN_LIB_DIR
     

   • 替换Your_GPU_Architecture为你的GPU架构（例如7.5对应V100）。

5. 编译PyTorch：

   • 使用make命令进行编译。这可能需要一些时间，具体取决于你的硬件性能：

     make -j$(nproc)
     

6. 安装PyTorch：

   • 编译完成后，运行以下命令来安装PyTorch：

     make install
     

验证安装

1. 验证安装：
   • 在Python解释器中导入PyTorch并检查版本：

     import torch
     print(torch.__version__)
     

通过以上步骤，你应该能够在Ubuntu上成功编译并安装PyTorch。如果在编译过程中遇到任何问题，请参考PyTorch的官方文档或GitHub仓库中的问题讨论。