PyTorch v2.1 + Caffe2 编译

源码编译 PyTorch 2.1 + Caffe2。

Caffe2 是当年大火的 Caffe 的升级，本来集成到了 PyTorch 官方仓库，但从 1.11.0 开始官方提供的预编译 wheel 就已经没有编译 Caffe2 了。目前部署依然依赖 detectron 和 Caffe2，需要自行编译。

系统环境：Ubuntu 20.04，Python 3.10，gcc 11.4，CUDA 12.1

  1. Clone PyTorch repo: git clone https://github.com/pytorch/pytorch.git
  2. 切换到 release/2.1: cd pytorch && git checkout release/2.1
  3. 更新子模块: git submodule update --init --recursive
  4. 修改 CMakeLists.txt: 最前面加上 set(CMAKE_CUDA_ARCHITECTURES "native") 和 set(CMAKE_CUDA_COMPILER "/usr/local/cuda/bin/nvcc")，后面 nvcc 的路径需要根据实际路径修改
  5. 编译: sudo CC=gcc CXX=g++ ARCHFLAGS="-arch x86_64" BUILD_CAFFE2=1 BUILD_CAFFE2_OPS=1 BUILD_PYTHON=1 python setup.py install
  6. 修改 PYTHONPATH 变量: export PYTHONPATH={path to pytorch/build}/$PYTHONPATH
  7. 测试 Caffe2 编译安装是否成功: python -c 'from caffe2.python import core' 2>/dev/null && echo "Success" || echo "Failure", 输出 Success 表示已成功编译安装

留言

2024-06-25
▲