PyTorch中apex的安装方式

酒酿小圆子～2024-12-18 21:02

apex是NVIDIA开发的基于PyTorch的混合精度训练加速神器，能够增加运算速度，并且减少显存的占用。

Github地址：https://github.com/NVIDIA/apex
官方教程：https://nvidia.github.io/apex/

安装方式

需要注意的是apex的安装不能通过 pip install 的方式安装（pip install apex安装的是两个完全不同的库），需要基于Github源码安装。官方提供的安装方式如下：

python 复制代码

git clone https://github.com/NVIDIA/apex
cd apex
# if pip >= 23.1 
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
# otherwise
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./

【注意】：Apex官方提供的下载命令很简单，但是安装过程中会有各种奇怪的报错。需要注意以下几点：

物理服务器上安装的CUDA版本
虚拟环境中安装的pytorch版本
虚拟环境中安装的cudatoolkit版本 （需要与服务器上CUDA版本一致！）

以上几个版本不匹配的话，就可能会报错。所以请务必确认好版本！

报错解决

由于apex兼容性太差，按照官网安装也可能会出现一些列问题，这时可以切换下载的Git分支，尝试哪个分支能够安装成功。具体可以参考：安装apex时遇到的问题