Mixed training precision leads to nan of loss


to save memory I am trying to train a resnet18 model in half precision. Therefore, I converted my inputs and my model. I also changed the BachNorm2d layers back to normal floating tensors with the following lines:

for layer in resnet18.modules():
    if isinstance(layer, nn.BatchNorm2d):

Nevertheless, I receive nan values as loss after the second batch is trained. Since, everything works find without half precision and the chance of instabilities during applying the BatchNorm2d, I was wondering if I am missing something or if there is another way?

Calling model.half() might work, but can easily create overflows and thus NaNs.
We recommend to use the automatic mixed-precision training, which takes care of these issues and also stabilizes the training using loss scaling.
You can use it in the nightly binaries or by building PyTorch from source.

1 Like