NaN's in gradients due to multi-objective loss function

username: Mahmoud_Abdelkhalek
cooked: <h2><a name="context-1" class="anchor" href="#context-1"></a>Context</h2>
<p>I’m currently trying to implement the following architecture:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/3/6312e07671f277807ed1e529664851756276bc76.png" data-download-href="https://discuss.pytorch.org/uploads/default/6312e07671f277807ed1e529664851756276bc76" title="Screenshot 2021-02-12 210143"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/6312e07671f277807ed1e529664851756276bc76_2_301x374.png" alt="Screenshot 2021-02-12 210143" data-base62-sha1="e8rIcznOpJIAMHz3ZmhEvVVH8uW" width="301" height="374" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/6312e07671f277807ed1e529664851756276bc76_2_301x374.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/6312e07671f277807ed1e529664851756276bc76_2_451x561.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/3/6312e07671f277807ed1e529664851756276bc76.png 2x" data-dominant-color="F3F3F3"><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">Screenshot 2021-02-12 210143</span><span class="informations">598×743 5.86 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>Where x is an audio signal, h_\theta is a linear filter, f_\phi is a network that predicts one of two classes, either y_hat = 0 or y_hat = 1, and g_\psi is an autoencoder that attempts to reconstruct x as x_\hat.</p>
<p>Here is how I pass a batch of x’s through the filter h_\theta:</p>
<pre data-code-wrap="python"><code class="lang-python">x = filt(x)
</code></pre>
<p>Where <code>filt</code> is a <code>nn.Module</code> object with trainable parameters:</p>
<pre data-code-wrap="python"><code class="lang-python">filt = Filter(bias = True, kernel_size = 201, identity = False).to(device)
</code></pre>
<p>After passing x through the filter h_\theta, here is how I pass it through the network f_\phi:</p>
<pre data-code-wrap="python"><code class="lang-python">def run_batch_disc(mode,x,labels,transforms,net,loss_func,device):
    
    log, mel, spec = transforms
    
    # compute logarithmic Mel-scale magnitude spectrograms of the filtered
    # signals
    
    x = spec(x)
    x = torchaudio.functional.complex_norm(x)
    x = mel(x)
    x = log(x)
    
    # with torch.set_grad_enabled(mode == 'train'):        
    
    # logits must have the same shape as labels
    
    logits = net(x).squeeze(dim = 1)
    
    # compute negative log-likelihood (NLL) using logits
    
    NLL = loss_func(logits,labels)
    
    # record predictions. since sigmoid(0) = 0.5, then negative values
    # correspond to class 0 and positive values correspond to class 1
    
    preds = logits &gt; 0
    
    # record correct predictions
    
    true_preds = torch.sum(preds == labels)
    
    return NLL,true_preds.item()
</code></pre>
<p>Where <code>log</code>, <code>mel</code>, and <code>spec</code> are:</p>
<pre data-code-wrap="python"><code class="lang-python">log = torchaudio.transforms.AmplitudeToDB(stype = 'magnitude',
                                          top_db = 80).to(device)

# settings for spectrograms

win_length_sec = 0.012
win_length = int(sample_rate * win_length_sec)
# need 50% overlap to satisfy constant-overlap-add constraint to allow
# for perfect reconstruction using inverse STFT
hop_length = int(sample_rate * win_length_sec / 2) # used to be 64
n_mels = 128
n_fft = 512 # used to be 1024

mel = torchaudio.transforms.MelScale(n_mels = n_mels,
                                     sample_rate = sample_rate,
                                     f_min = 0.0,
                                     f_max = None,
                                     n_stft = n_fft // 2 + 1).to(device)

spec = torchaudio.transforms.Spectrogram(n_fft = n_fft,
                                         win_length = win_length,
                                         hop_length = hop_length,
                                         pad = 0,
                                         window_fn = torch.hann_window,
                                         power = None,
                                         normalized = False,
                                         wkwargs = None).to(device)
</code></pre>
<p>In parallel, I am pass the filtered x through the autoencoder g_\psi as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">def run_batch_recon(mode,x,spec,net,loss_func,device):
    
    # compute magnitude spectrogram of the filtered signal
    
    x = torchaudio.functional.complex_norm(spec(x))
    
    # scale each magnitude spectrogram in the batch to the interval [0,1]
    
    scale_factor = x.amax(dim=(2,3))[(..., ) + (None, ) * 2]
    x_scaled = x / scale_factor
    
    # with torch.set_grad_enabled(mode == 'train'):
        
    # compute reconstruction of the scaled magnitude spectrogram
    
    x_hat = net(x_scaled)
    
    # unscale magnitude spectrogram back to normal values
    
    x_hat = x_hat * scale_factor
    
    # compute reconstruction loss
    
    recon_loss = loss_func(x_hat,x)
    
    return recon_loss
</code></pre>
<p>Where <code>loss_func</code> is:</p>
<pre data-code-wrap="python"><code class="lang-python">def recon_loss_func(x_hat,x,alpha = 1):
    
    # spectral convergence
    
    num = (x - x_hat).pow(2).sum(dim=(2,3)).sqrt().squeeze()
    den = x.pow(2).sum(dim=(2,3)).sqrt().squeeze()
    spec_conv = torch.div(num,den)
    
    # log-scale STFT magnitude loss
    
    eps = 1e-10
    num = (torch.log(x + eps) - torch.log(x_hat + eps)).abs().sum(dim=2).amax(dim=2).squeeze()
    den = torch.log(x + eps).abs().sum(dim=2).amax(dim=2).squeeze()
    log_loss = torch.div(num,den)
    
    return torch.mean(spec_conv + alpha * log_loss)
</code></pre>
<h2><a name="objective-2" class="anchor" href="#objective-2"></a>Objective</h2>
<p>I am trying to choose the parameters \theta, \phi, and \psi to minimize the following objective function:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/a/aac91516b5d5001a8c4c8aecf6002c2533203024.png" data-download-href="https://discuss.pytorch.org/uploads/default/aac91516b5d5001a8c4c8aecf6002c2533203024" title="loss_func"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/a/aac91516b5d5001a8c4c8aecf6002c2533203024_2_517x57.png" alt="loss_func" data-base62-sha1="omPZbXSGjqqQL6pTwxlPXDdaigs" width="517" height="57" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/a/aac91516b5d5001a8c4c8aecf6002c2533203024_2_517x57.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/a/aac91516b5d5001a8c4c8aecf6002c2533203024_2_775x85.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/a/aac91516b5d5001a8c4c8aecf6002c2533203024_2_1034x114.png 2x" data-dominant-color="F1F1F1"><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">loss_func</span><span class="informations">1340×150 8.85 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>Where L_y is the binary cross entropy loss of predicting y_hat not equal to y (<code>NLL</code> in the code above) and L_x is a modified RMSE loss between x and its reconstruction x_hat (<code>recon_loss_func</code> above). Essentially, I am trying to minimize the binary cross entropy loss while maximizing the MSE reconstruction loss. The purpose of this is to choose the filter parameters \theta that balance these two objectives.</p>
<p>Here is how I implement this objective function:</p>
<pre data-code-wrap="python"><code class="lang-python">def run_batch(batch,mode,transforms,filt,nets,loss_funcs,lambda_adv,
              optimizer,device):
    
    # unpack everything
    
    x,labels = batch
    disc_net,recon_net = nets
    disc_loss_func, recon_loss_func = loss_funcs
    
    # move to GPU if available
    
    x = x.to(device)
    labels = labels.to(device).type_as(x) # needed for NLL
    
    with torch.set_grad_enabled(mode == 'train'):
        
        # filter each signal
        
        x = filt(x)
        
        # compute classification (negative log-likelihood (NLL)) loss and
        # number of correct predictions
        
        NLL, true_preds = run_batch_disc(mode,x,labels,transforms,disc_net,
                                         disc_loss_func,device)
        
        # extract speech signals from x and compute reconstruction loss
        
        where_speech = torch.nonzero(labels == 0, as_tuple = True)
        recon_loss = run_batch_recon(mode,x[where_speech],transforms[2],
                                     recon_net,recon_loss_func,device)
        
        # compute adversarial loss
        
        adv_loss = NLL - lambda_adv * recon_loss
        
        if mode == 'train':
        
            # compute gradient of adversarial loss with respect to filter,
            # discriminator network, and reconstruction network parameters
            
            adv_loss.backward()
            
            # update parameters using the gradient descent update rule
            
            optimizer.step()
            
            # zero the accumulated parameter gradients
            
            optimizer.zero_grad()
    
    return NLL.item(),true_preds,recon_loss.item(),adv_loss.item()
</code></pre>
<h2><a name="the-problem-3" class="anchor" href="#the-problem-3"></a>The Problem</h2>
<p>As I am trying to implement this, I keep getting all NaN’s in the gradients of the filter parameters \theta once I call <code>.backward()</code> on the objective function shown above, while the gradients of the \phi and \psi of f and g respectively do not contain NaN’s.</p>
<p>I have tried to set</p>
<pre data-code-wrap="python"><code class="lang-python">torch.autograd.set_detect_anomaly(True)
</code></pre>
<p>And I get the following <code>RuntimeError</code> and traceback:</p>
<pre data-code-wrap="python"><code class="lang-python">RuntimeError: Function 'PowBackward0' returned nan values in its 0th output.
C:\ProgramData\Anaconda3\lib\site-packages\torch\autograd\__init__.py:130: UserWarning: Error detected in PowBackward0. Traceback of forward call that caused the error:
  File "c:\users\mahmoud talaat\documents\ncsu\semesters\fall2020\ece695_masters_research\adversarial_train_val.py", line 323, in &lt;module&gt;
    train_stats = run_epoch('train',
  File "c:\users\mahmoud talaat\documents\ncsu\semesters\fall2020\ece695_masters_research\adversarial_train_val.py", line 151, in run_epoch
    stats = run_batch(batch,mode,transforms,filt,nets,loss_funcs,
  File "c:\users\mahmoud talaat\documents\ncsu\semesters\fall2020\ece695_masters_research\adversarial_train_val.py", line 94, in run_batch
    NLL, true_preds = run_batch_disc(mode,x,labels,transforms,disc_net,
  File "c:\users\mahmoud talaat\documents\ncsu\semesters\fall2020\ece695_masters_research\adversarial_train_val.py", line 19, in run_batch_disc
    x = torchaudio.functional.complex_norm(x)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torchaudio\functional.py", line 424, in complex_norm
    return complex_tensor.pow(2.).sum(-1).pow(0.5 * power)
 (Triggered internally at  ..\torch\csrc\autograd\python_anomaly_mode.cpp:104.)
  Variable._execution_engine.run_backward(
</code></pre>
<p>The part where this <code>RuntimeError</code> happens is here:</p>
<pre data-code-wrap="python"><code class="lang-python">def run_batch_disc(mode,x,labels,transforms,net,loss_func,device):
    
    log, mel, spec = transforms
    
    # compute logarithmic Mel-scale magnitude spectrograms of the filtered
    # signals
    
    x = spec(x)
    x = torchaudio.functional.complex_norm(x)
    x = mel(x)
    x = log(x)
</code></pre>
<p>I have made sure that <code>x = spec(x)</code> is &gt;= 0, but I still get this problem.</p>
<p>I have noticed that if I remove the autoencoder g_\psi, and just pass x through h and f and call backward on the binary cross entropy loss alone, there are no NaN’s in the gradients of the filter parameters.</p>

username: ptrblck
cooked: <p>Could you make sure that the <code>log</code> operation returns valid outputs, e.g. that <code>x</code> is not zero?<br>
Given that <code>spec(x)</code> can be zero, would it also be possible that <code>complex_norm</code> and <code>mel</code> return zeros?<br>
If so, then <code>torch.log(torch.zeros(1))</code> would return <code>tensor([-inf])</code>, which might raise this issue in the backward pass.</p>

username: Mahmoud_Abdelkhalek
cooked: <p>This solved the problem!</p>
<p>Just needed to change this:</p>
<pre><code class="lang-python">x = spec(x)
x = torchaudio.functional.complex_norm(x)
x = mel(x)
x = log(x)
</code></pre>
<p>To this:</p>
<pre><code class="lang-python">eps = 1e-9
x = spec(x)
x = torchaudio.functional.complex_norm(x + eps)
x = mel(x)
x = log(x)
</code></pre>
<p>I wasn’t sure if <code>log(x)</code> will cause a problem if <code>x</code> contains zeroes, so I added <code>eps</code> to it as well:</p>
<pre><code class="lang-python">eps = 1e-9
x = spec(x)
x = torchaudio.functional.complex_norm(x + eps)
x = mel(x)
x = log(x + eps)
</code></pre>