README.md 4.73 KB
Newer Older
Christian Sarofeen's avatar
Christian Sarofeen committed
1
2
# Introduction

Michael Carilli's avatar
Michael Carilli committed
3
4
5
6
7
8
This repository holds NVIDIA-maintained utilities to streamline 
mixed precision and distributed training in Pytorch. 
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.

9
## Full API Documentation: [https://nvidia.github.io/apex](https://nvidia.github.io/apex)
Michael Carilli's avatar
Michael Carilli committed
10
11
12

# Contents

mcarilli's avatar
mcarilli committed
13
## 1. Amp:  Automatic Mixed Precision
Michael Carilli's avatar
Michael Carilli committed
14

Michael Carilli's avatar
Michael Carilli committed
15
16
17
`apex.amp` is a tool to enable mixed precision training by changing only 3 lines of your script.
Users can easily experiment with different pure and mixed precision training modes by supplying
different flags to `amp.initialize`.
Michael Carilli's avatar
Michael Carilli committed
18

Michael Carilli's avatar
Michael Carilli committed
19
20
[Webinar introducing Amp](https://info.nvidia.com/webinar-mixed-precision-with-pytorch-reg-page.html)
(The flag `cast_batchnorm` has been renamed to `keep_batchnorm_fp32`).
Michael Carilli's avatar
Michael Carilli committed
21

Michael Carilli's avatar
Michael Carilli committed
22
[API Documentation](https://nvidia.github.io/apex/amp.html)
Michael Carilli's avatar
Michael Carilli committed
23

Michael Carilli's avatar
Michael Carilli committed
24
[Comprehensive Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
Michael Carilli's avatar
Michael Carilli committed
25

Michael Carilli's avatar
Michael Carilli committed
26
[DCGAN example coming soon...](https://github.com/NVIDIA/apex/tree/master/examples/dcgan)
27

28
[Moving to the new Amp API](https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users) (for users of the deprecated "Amp" and "FP16_Optimizer" APIs)
Michael Carilli's avatar
Michael Carilli committed
29
30
31
32
33
34
35

## 2. Distributed Training

`apex.parallel.DistributedDataParallel` is a module wrapper, similar to 
`torch.nn.parallel.DistributedDataParallel`.  It enables convenient multiprocess distributed training,
optimized for NVIDIA's NCCL communication library.

36
37
[API Documentation](https://nvidia.github.io/apex/parallel.html)

38
[Python Source](https://github.com/NVIDIA/apex/tree/master/apex/parallel)
Michael Carilli's avatar
Michael Carilli committed
39

40
[Example/Walkthrough](https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed)
Christian Sarofeen's avatar
Christian Sarofeen committed
41

Michael Carilli's avatar
Michael Carilli committed
42
43
The [Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
shows use of `apex.parallel.DistributedDataParallel` along with `apex.amp`.
44

jjsjann123's avatar
jjsjann123 committed
45
46
47
48
### Synchronized Batch Normalization

`apex.parallel.SyncBatchNorm` extends `torch.nn.modules.batchnorm._BatchNorm` to
support synchronized BN.
Michael Carilli's avatar
Michael Carilli committed
49
50
51
52
53
54
55
It allreduces stats across processes during multiprocess (DistributedDataParallel) training.
Synchronous BN has been used in cases where only a small
local minibatch can fit on each GPU.
Allreduced stats increase the effective batch size for the BN layer to the
global batch size across all processes (which, technically, is the correct
formulation).
Synchronous BN has been observed to improve converged accuracy in some of our research models.
jjsjann123's avatar
jjsjann123 committed
56

Christian Sarofeen's avatar
Christian Sarofeen committed
57
58
59
# Requirements

Python 3
Michael Carilli's avatar
Michael Carilli committed
60

Michael Carilli's avatar
Michael Carilli committed
61
CUDA 9 or newer
Michael Carilli's avatar
Michael Carilli committed
62

Michael Carilli's avatar
Michael Carilli committed
63
PyTorch 0.4 or newer.  The CUDA and C++ extensions require pytorch 1.0 or newer.
mcarilli's avatar
mcarilli committed
64

Michael Carilli's avatar
Michael Carilli committed
65
66
We recommend the latest stable release, obtainable from
[https://pytorch.org/](https://pytorch.org/).  We also test against the latest master branch, obtainable from [https://github.com/pytorch/pytorch](https://github.com/pytorch/pytorch).
Christian Sarofeen's avatar
Christian Sarofeen committed
67

Michael Carilli's avatar
Michael Carilli committed
68
69
70
It's often convenient to use Apex in Docker containers.  Compatible options include:
* [NVIDIA Pytorch containers from NGC](https://ngc.nvidia.com/catalog/containers/nvidia%2Fpytorch), which come with Apex preinstalled.  To use the latest Amp API, you may need to `pip uninstall apex` then reinstall Apex using the **Quick Start** commands below.
* [official Pytorch -devel Dockerfiles](https://hub.docker.com/r/pytorch/pytorch/tags), e.g. `docker pull pytorch/pytorch:nightly-devel-cuda10.0-cudnn7`, in which you can install Apex using the **Quick Start** commands.
Christian Sarofeen's avatar
Christian Sarofeen committed
71

72
73
See the [Docker example folder](https://github.com/NVIDIA/apex/tree/master/examples/docker) for details.

Christian Sarofeen's avatar
Christian Sarofeen committed
74
75
# Quick Start

76
### Linux
Christian Sarofeen's avatar
Christian Sarofeen committed
77

Michael Carilli's avatar
Michael Carilli committed
78
79
For performance and full functionality, we recommend installing Apex with
CUDA and C++ extensions via
Christian Sarofeen's avatar
Christian Sarofeen committed
80
```
Glenn Jocher's avatar
Glenn Jocher committed
81
$ git clone https://github.com/NVIDIA/apex
Michael Carilli's avatar
Michael Carilli committed
82
83
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
Christian Sarofeen's avatar
Christian Sarofeen committed
84
85
```

Michael Carilli's avatar
Michael Carilli committed
86
Apex also supports a Python-only build (required with Pytorch 0.4) via
jjsjann123's avatar
jjsjann123 committed
87
```
Michael Carilli's avatar
Michael Carilli committed
88
$ pip install -v --no-cache-dir .
jjsjann123's avatar
jjsjann123 committed
89
```
Michael Carilli's avatar
Michael Carilli committed
90
A Python-only build omits:
91
- Fused kernels required to use `apex.optimizers.FusedAdam`.
vfdev's avatar
typo  
vfdev committed
92
- Fused kernels required to use `apex.normalization.FusedLayerNorm`.
Michael Carilli's avatar
Michael Carilli committed
93
94
95
- Fused kernels that improve the performance and numerical stability of `apex.parallel.SyncBatchNorm`.
- Fused kernels that improve the performance of `apex.parallel.DistributedDataParallel` and `apex.amp`.
`DistributedDataParallel`, `amp`, and `SyncBatchNorm` will still be usable, but they may be slower.
jjsjann123's avatar
jjsjann123 committed
96

97
### Windows support
mcarilli's avatar
mcarilli committed
98
99
Windows support is experimental, and Linux is recommended.  `pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .` may work if you were able to build Pytorch from source
on your system.  `pip install -v --no-cache-dir .` (without CUDA/C++ extensions) is more likely to work.  If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.