This repo is designed to hold PyTorch modules and utilities that are under active development and experimental. This repo is not designed as a long term solution or a production solution. Things placed in here are intended to be eventually moved to upstream PyTorch.
Python 3 PyTorch 0.3 or newer CUDA 9
To build the extension run the following command in the root directory of this project
python setup.py install
To use the extension simply run
import apex
and optionally (if required for your use)
import apex._C as apex_backend
Current version of apex contains:
- Mixed precision utilities can be found here examples of using mixed precision utilities can be found for the PyTorch imagenet example and the PyTorch word language model example.
- Parallel utilities can be found here and an example/walkthrough can be found here
- apex/parallel/distributed.py contains a simplified implementation of PyTorch's DistributedDataParallel that's optimized for use with NCCL in single gpu / process mode
- apex/parallel/multiproc.py is a simple multi-process launcher that can be used on a single node/computer with multiple GPU's
- Reparameterization function that allows you to recursively apply reparameterization to an entire module (including children modules).
- An experimental and in development flexible RNN API.