Model learns nothing with mixed precision training

username: jbohnslav
cooked: <p>I am currently using PyTorch for self-supervised optic flow training. The new AMP package is very appealing, as it the model runs faster and can use a larger batch size when using FP16 activations. However, I’ve found that the model does not learn anything.</p>
<p>Here is an ugly image, with FP32 on the left and mixed precision on the right. I used identical hyperparameters, and ran for longer than necessary to prove the point.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/8/082f2c979207e698a01d7d957fe85ebe401f7535.png" data-download-href="https://discuss.pytorch.org/uploads/default/082f2c979207e698a01d7d957fe85ebe401f7535" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/8/082f2c979207e698a01d7d957fe85ebe401f7535.png" alt="image" data-base62-sha1="1aoTb8FWUd5DZQpvkmRmKUiL0W1" width="690" height="219" data-dominant-color="E4EAE4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">697×222 29.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The relevant code was taken straight from the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html" rel="noopener nofollow ugc">amp example page</a>. Here is the relevant portion copied:</p>
<pre data-code-wrap="python"><code class="lang-python">if fp16:
    with autocast():
        outputs = model(batch)
        downsampled_t0, estimated_t0, flows_reshaped = reconstructor(batch, outputs)
        loss, loss_components = criterion(batch, downsampled_t0, estimated_t0, flows_reshaped)
else:
    outputs = model(batch)
    downsampled_t0, estimated_t0, flows_reshaped = reconstructor(batch, outputs)
    loss, loss_components = criterion(batch, downsampled_t0, estimated_t0, flows_reshaped)

# zero the parameter gradients
optimizer.zero_grad()
# calculate gradients
if fp16:
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
else:
    loss.backward()
    # step in direction of gradients according to optimizer
    optimizer.step()
</code></pre>
<p>Things that could be causing the issue</p>
<ul>
<li>Using a multi-part loss function. I have multiple image reconstruction losses and a smoothness loss, common in optic flow tasks. The AMP page tells us to scale components of our loss function individually. However, this seems to make little sense to me. The relevant portions of my loss function have been carefully scaled by prior work and my own work, so scaling them individually seems like it would ruin that. Furthermore, having a backwards pass for each component of my loss function would take much longer.</li>
<li>Spatial transformer networks. I sample from the image at time t+1 to estimate the image at time t using torch.nn.functional.grid_sample. Does this operation not work with mixed precision training?</li>
<li>Just a bug in torch somewhere (I can open a GitHub issue if so)</li>
</ul>
<p>Debug info:</p>
<ul>
<li>pytorch version: 1.7.0a0+eb47940</li>
<li>cuda: 11</li>
</ul>

username: ptrblck
cooked: <aside class="quote no-group" data-username="jbohnslav" data-post="1" data-topic="93420">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/958977/48.png" class="avatar"> jbohnslav:</div>
<blockquote>
<p>However, this seems to make little sense to me. The relevant portions of my loss function have been carefully scaled by prior work and my own work, so scaling them individually seems like it would ruin that. Furthermore, having a backwards pass for each component of my loss function would take much longer.</p>
</blockquote>
</aside>
<p>The <code>GradScaler</code> will scale the loss to increase the magnitude of the gradients and will unscale it using the same scaling factor. Your manual scaling won’t be changed, so I would recommend to try it out and check, if it solves the issue.</p>