Apex (A PyTorch Extension)¶
This site contains the API documentation for Apex (https://github.com/nvidia/apex), a Pytorch extension with NVIDIA-maintained utilities to streamline mixed precision and distributed training. Some of the code here will be included in upstream Pytorch eventually. The intention of Apex is to make up-to-date utilities available to users as quickly as possible.
Installation requires CUDA 9 or later, PyTorch 0.4 or later, and Python 3. Install by running
git clone https://www.github.com/nvidia/apex
cd apex
python setup.py install [--cuda_ext] [--cpp_ext]
AMP: Automatic Mixed Precision
FP16/Mixed Precision Utilities
Distributed Training
Fused Optimizers
Fused Layer Norm