Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

Mixed Precision DCGAN Training in PyTorch

main_amp.py is based on https://github.com/pytorch/examples/tree/master/dcgan. It implements Automatic Mixed Precision (Amp) training of the DCGAN example for different datasets. Command-line flags forwarded to amp.initialize are used to easily manipulate and switch between various pure and mixed precision "optimization levels" or opt_levels. For a detailed explanation of opt_levels, see the updated API guide.

We introduce these changes to the PyTorch DCGAN example as described in the Multiple models/optimizers/losses section of the documentation::

# Added after models and optimizers construction
[netD, netG], [optimizerD, optimizerG] = amp.initialize(
    [netD, netG], [optimizerD, optimizerG], opt_level=opt.opt_level, num_losses=3)
...
# loss.backward() changed to:
with amp.scale_loss(errD_real, optimizerD, loss_id=0) as errD_real_scaled:
    errD_real_scaled.backward()
...
with amp.scale_loss(errD_fake, optimizerD, loss_id=1) as errD_fake_scaled:
    errD_fake_scaled.backward()
...
with amp.scale_loss(errG, optimizerG, loss_id=2) as errG_scaled:
    errG_scaled.backward()

Note that we use different loss_scalers for each computed loss. Using a separate loss scaler per loss is optional, not required.

To improve the numerical stability, we swapped nn.Sigmoid() + nn.BCELoss() to nn.BCEWithLogitsLoss().

With the new Amp API you never need to explicitly convert your model, or the input data, to half().

"Pure FP32" training:

$ python main_amp.py --opt_level O0

Recommended mixed precision training:

$ python main_amp.py --opt_level O1

Have a look at the original DCGAN example for more information about the used arguments.

To enable mixed precision training, we introduce the --opt_level argument.