如何在Ubuntu上编译PyTorch源码

ubuntu

小樊

2025-11-29 01:42:31

栏目: 智能运维

在Ubuntu上编译PyTorch源码需要遵循以下步骤：

安装依赖项：在开始编译之前，确保你的系统已经安装了所有必要的依赖项。这通常包括：
- Python开发工具和库（如python3, python3-dev, pip）
- CUDA和cuDNN（如果你打算使用GPU加速）
- CMake
- Ninja构建系统
- BLAS库（如OpenBLAS或MKL）
- Git
你可以使用以下命令来安装这些依赖项：
```
sudo apt-get update
sudo apt-get install -y python3 python3-dev python3-pip cmake ninja-build git
# 如果使用GPU，安装CUDA和cuDNN
```
克隆PyTorch仓库：使用Git克隆PyTorch的GitHub仓库：
```
git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
```
--recursive标志用于克隆所有子模块。

设置环境变量：如果你打算使用CUDA，确保设置了CUDA_HOME环境变量，并将其添加到你的PATH中：

export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

你还需要设置CMAKE_PREFIX_PATH，以便PyTorch可以找到BLAS库：

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

安装Python依赖项：使用pip安装PyTorch所需的Python包：
```
pip install -r requirements.txt
```
构建PyTorch：使用CMake和Ninja来构建PyTorch。你可以使用以下命令来配置和构建项目：
```
python setup.py install
```
或者，如果你想使用Ninja构建系统，可以运行：
```
export USE_NINJA=1
python setup.py install
```
验证安装：安装完成后，你可以通过运行一些PyTorch测试来验证安装是否成功：
```
python -c "import torch; print(torch.__version__)"
```
这应该会打印出PyTorch的版本号。

请注意，编译PyTorch源码可能需要一些时间，具体取决于你的硬件性能。如果你遇到任何问题，可以查看PyTorch的官方文档或在社区寻求帮助。

最新问答