Package: pytorch-extension
IntroductionThis is a Python package available onPyPIfor 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.Full API Documentation:https://nvidia.github.io/apexGTC 2019andPytorch DevCon 2019SlidesContents1. Amp:  Automatic Mixed Precisionapex.ampis 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 toamp.initialize.Webinar introducing Amp(The flagcast_batchnormhas been renamed tokeep_batchnorm_fp32).API DocumentationComprehensive Imagenet exampleDCGAN example coming soon...Moving to the new Amp API(for users of the deprecated "Amp" and "FP16_Optimizer" APIs)2. Distributed Trainingapex.parallel.DistributedDataParallelis a module wrapper, similar totorch.nn.parallel.DistributedDataParallel.  It enables convenient multiprocess distributed training,
optimized for NVIDIA's NCCL communication library.API DocumentationPython SourceExample/WalkthroughTheImagenet exampleshows use ofapex.parallel.DistributedDataParallelalong withapex.amp.Synchronized Batch Normalizationapex.parallel.SyncBatchNormextendstorch.nn.modules.batchnorm._BatchNormto
support synchronized BN.
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.CheckpointingTo properly save and load youramptraining, we introduce theamp.state_dict(), which contains allloss_scalersand their corresponding unskipped steps,
as well asamp.load_state_dict()to restore these attributes.In order to get bitwise accuracy, we recommend the following workflow:# Initializationopt_level='O1'model,optimizer=amp.initialize(model,optimizer,opt_level=opt_level)# Train your model...withamp.scale_loss(loss,optimizer)asscaled_loss:scaled_loss.backward()...# Save checkpointcheckpoint={'model':model.state_dict(),'optimizer':optimizer.state_dict(),'amp':amp.state_dict()}torch.save(checkpoint,'amp_checkpoint.pt')...# Restoremodel=...optimizer=...checkpoint=torch.load('amp_checkpoint.pt')model,optimizer=amp.initialize(model,optimizer,opt_level=opt_level)model.load_state_dict(checkpoint['model'])optimizer.load_state_dict(checkpoint['optimizer'])amp.load_state_dict(checkpoint['amp'])# Continue training...Note that we recommend restoring the model using the sameopt_level. Also note that we recommend calling theload_state_dictmethods afteramp.initialize.RequirementsPython 3CUDA 9 or newerPyTorch 0.4 or newer.  The CUDA and C++ extensions require pytorch 1.0 or newer.Quick StartLinuxFor performance and full functionality, we recommend installing with CUDA and C++ extensions according topip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" pytorch-extensionFor a Python-only build (required with Pytorch 0.4):pip install -v --disable-pip-version-check --no-cache-dir pytorch-extensionA Python-only build omits:Fused kernels required to useapex.optimizers.FusedAdam.Fused kernels required to useapex.normalization.FusedLayerNorm.Fused kernels that improve the performance and numerical stability ofapex.parallel.SyncBatchNorm.Fused kernels that improve the performance ofapex.parallel.DistributedDataParallelandapex.amp.DistributedDataParallel,amp, andSyncBatchNormwill still be usable, but they may be slower.Pyprof support has been moved to its owndedicated repository.
The codebase is deprecated in Apex and will be removed soon.Windows supportWindows support is experimental, and Linux is recommended.pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" pytorch-extensionmay work if you were able to build Pytorch from source
on your system.pip install -v --disable-pip-version-check --no-cache-dir pytorch-extension(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.