Help with autocast causing loss issue

username: musl
cooked: <p>tl;dr Using autocast is causing issues with a specific loss. Forcing to fp32 didn’t fix. Help needed.</p>
<p><strong>Problem description:</strong></p>
<p>Autocast and Gradscale is being <a href="https://github.com/muslll/neosr/blob/master/neosr/models/default.py#L230" rel="noopener nofollow ugc">applied</a> on the training loop, as AMP documentation explains:</p>
<pre><code class="lang-auto">        self.optimizer_g.zero_grad(set_to_none=True)
        scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

        with torch.autocast(device_type='cuda', dtype=amp_dtype, enabled=use_amp):
            self.output = self.net_g(self.lq)
            l_g_total = 0
            loss_dict = OrderedDict()
            if self.cri_perceptual:
                l_g_percep, l_g_style = self.cri_perceptual(self.output, self.gt)
                if l_g_percep is not None:
                    l_g_total += l_g_percep
                    loss_dict['l_percep'] = l_g_percep

        scaler.scale(l_g_total).backward()
        scaler.step(self.optimizer_g)
        scaler.update()
</code></pre>
<p>The <code>cri_perceptual</code> is the <a href="https://github.com/muslll/neosr/blob/master/neosr/losses/basic_loss.py#L173" rel="noopener nofollow ugc">Perceptual Loss</a>, which is causing issues:</p>
<pre><code class="lang-auto">def forward(self, x, gt):
	x_features = self.vgg(x)
	gt_features = self.vgg(gt.detach())
	for k in x_features.keys():
		percep_loss += self.criterion(
                        x_features[k], gt_features[k]) * self.layer_weights[k]
	return percep_loss
</code></pre>
<p>It uses torchvision VGG model to extract features:</p>
<pre><code class="lang-auto">from torchvision.models import vgg
from torchvision.models import VGG19_Weights

vgg_net = getattr(vgg, vgg_type)(weights=VGG19_Weights.DEFAULT)
self.vgg_net.eval()

def forward(self, x):
	if self.range_norm:
            x = (x + 1) / 2
        if self.use_input_norm:
            x = (x - self.mean) / self.std

        output = {}
        for key, layer in self.vgg_net._modules.items():
            x = layer(x)
            if key in self.layer_name_list:
                output[key] = x.clone()

        return output
</code></pre>
<p><strong>Things I’ve tried so far, without success:</strong></p>
<ul>
<li>Adding decorator <code>@torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)</code> to both vgg feature extraction and perceptual forward. Didn’t work, even though I did not use <code>custom_bwd</code> as documentation says, as this would require some modifications on current <code>.backward()</code></li>
<li>Forcing VGG inference with <code>.float()</code> or <code>.half()</code></li>
</ul>
<p><strong>Question:</strong></p>
<ul>
<li>Any idea how to solve this and what might be causing the issue?</li>
</ul>

username: musl
cooked: <p>Update: I was able to identify the origin of the issue is in GradScaler backward().<br>
Still not sure how to solve it.</p>

username: musl
cooked: <p>Solution has been found by one of my projects users (<a href="https://github.com/terrainer" rel="noopener nofollow ugc">@terrainer</a>).<br>
The initial GradScaler value was too high, leading to inf/NaNs and making it unable to recover even after <code>backoff_factor</code> being applied. Decreasing <code>init_scale</code> from default <code>2.**16</code> to <code>2.**11</code> fixes the issue. Thanks again to <span class="mention">@terrainer</span> for finding, testing and reporting the solution.</p>