Custom loss function doesn't update weights VGG-16

pvardanis: <p>Before I start, I just want you to know that I’ve read all the previous threads regarding this but still my problem persists. So, I’ve implemented a custom loss function that looks like this:</p>
<pre><code class="lang-auto">def Cosine(output, target):
    '''
    Custom loss function with 2 losses:

    - loss_1: penalizes the area out of the unit circle 
    - loss_2: 0 if output = target

    Inputs
        output: predicted phases
        target: true phases  
    ''' 

    # Penalize if output is out of the unit circle   
    squares = output ** 2 # (x ^ 2, y ^ 2)
    loss_1 = ((squares[:, ::2] + squares[:, 1::2]) - 1) ** 2 # (x ^ 2 + y ^ 2 - 1) ** 2

    # Compute the second loss, 1 - cos
    loss_2 =  1. - torch.cos(torch.atan2(output[:, 1::2], output[:, ::2]) - target)  
    
    return torch.mean(loss_1 + loss_2)
</code></pre>
<p>I’ve tried the following to debug my model:</p>
<pre><code class="lang-auto">for epoch in range(self.num_epochs):
                # Train
                network.train() # keep grads
                print('\nEpoch {}'.format(epoch+1))
                print('\nTrain:\n')
                a = list(network.parameters())[0].clone()
                for batch_idx, (images, labels) in enumerate(Bar(loaders['train'])):
                    images, labels = images.to(self.device, dtype=torch.float), labels.to(self.device, dtype=torch.float) # labels is a tensor of (512, 128) values if we use MyVgg
                    optimizer.zero_grad()
                    preds = network(images) 
                    loss = self.criterion(preds, labels)
                    loss.backward()
                    optimizer.step()
                    
                b = list(network.parameters())[0].clone()     
                print(torch.equal(a.data, b.data))  # This prints always true 
                
                # Validation
                print('\nValid:\n')
                network.eval() # skips dropout and batch_norm 
                with torch.no_grad():
                    for batch_idx, (images, labels) in enumerate(Bar(loaders['valid'])):
                        images, labels = images.to(self.device, dtype=torch.float), labels.to(self.device, 
                        dtype=torch.float) 
                        preds = network(images)
                        loss = self.criterion(preds, labels)

                if self.lr_scheduler: 
                    scheduler.step(loss) # update lr_scheduler
</code></pre>
<p>Wherever I try to print <code>list(self.network.parameters())[0].grad is None)</code> I always get <strong>True</strong>. Also, <code>model.parameters()</code> are always the same and don’t change. Model prediction, true labels, loss functions are all as I expected, but the backprop doesn’t seem to work.</p>
<p>Printing the loss gives me this:</p>
<pre><code class="lang-auto">tensor(1.0463, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
</code></pre>

KFrank: <p>Hello pvardanis!</p>
<aside class="quote no-group" data-username="pvardanis" data-post="1" data-topic="72944" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/pvardanis/48/21308_2.png" class="avatar"> pvardanis:</div>
<blockquote>
<p>Wherever I try to print <code>list(self.network.parameters())[0].grad is None)</code> I always get <strong>True</strong>.</p>
</blockquote>
</aside>
<p>From your other threads I have the impression that you are using<br>
a pre-trained VGG-16 network with a little of your own network<br>
architecture on the back end.  (Is this correct?)</p>
<p>Is it possible that you are freezing the VGG-16 weights (not<br>
unreasonable) and only updating (“fine-tuning”) the weights of<br>
your back-end architecture?  If so, it is possible that<br>
<code>list(network.parameters())[0]</code> is a frozen VGG-16<br>
weight, so that it doesn’t have a <code>grad</code>, and  isn’t changed by<br>
your optimization step?</p>
<p>Best.</p>
<p>K. Frank</p>

pvardanis: <p>Hi KFrank,</p>
<p>You are correct. I guess that’s what happens when you’re new to PyTorch haha. I printed the wrong grads. So, I managed to have some plots for all gradients on the classifier, with tensorboard. Again, the model I used is VGG-16 with batch norm and already pre-trained weights. I freezed all the feature layers, and trained the layers on the classifier. And by seeing the values I’m getting suspicious.</p>
<p>Hyperparameters used for this one:</p>
<ul>
<li>num_epochs = 10</li>
<li>batch_size = 64</li>
<li>lr = 0.01</li>
<li>SGD optimizer with default settings</li>
</ul>
<p>I noticed that if I lower the batch size to 16 or less I get NaN values as output.</p>
<p>Here are the gradients <strong>histograms</strong> for the <strong>3 layers of the classifier</strong> (0, 3 and 5, the layers in between are dropout and relu that’s why they’re not visualized):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/d/dd1af0b2d930031ce6ae3cbdff30060411aca8bb.png" data-download-href="https://discuss.pytorch.org/uploads/default/dd1af0b2d930031ce6ae3cbdff30060411aca8bb" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/d/dd1af0b2d930031ce6ae3cbdff30060411aca8bb_2_690x485.png" alt="image" data-base62-sha1="vxZeT6XWzaiaVb052L1UPQTpwlZ" width="690" height="485" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/d/dd1af0b2d930031ce6ae3cbdff30060411aca8bb_2_690x485.png, https://discuss.pytorch.org/uploads/default/original/3X/d/d/dd1af0b2d930031ce6ae3cbdff30060411aca8bb.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/d/dd1af0b2d930031ce6ae3cbdff30060411aca8bb.png 2x" data-dominant-color="FBF2EF"><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">850×598 44.5 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>And here are the gradients <strong>distributions</strong> for the <strong>3 layers of the classifier</strong>:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/6/f65088f87ab47e0ce810771fe2fbeb92a1cd6f63.png" data-download-href="https://discuss.pytorch.org/uploads/default/f65088f87ab47e0ce810771fe2fbeb92a1cd6f63" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/6/f65088f87ab47e0ce810771fe2fbeb92a1cd6f63_2_663x500.png" alt="image" data-base62-sha1="z900cxd2vdg4VsmbJwKsiEaO9vJ" width="663" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/6/f65088f87ab47e0ce810771fe2fbeb92a1cd6f63_2_663x500.png, https://discuss.pytorch.org/uploads/default/original/3X/f/6/f65088f87ab47e0ce810771fe2fbeb92a1cd6f63.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/6/f65088f87ab47e0ce810771fe2fbeb92a1cd6f63.png 2x" data-dominant-color="FDF5F2"><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">795×599 40.3 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>This is pretty damn low, or I’m wrong? Seems like a vanishing gradient problem to me, but I’m not sure.</p>

pvardanis: <p>I tried also the following approach:</p>
<p>Tanh activation function, with 1024 neurons at the end. Because of tanh each value is in the range [-1, 1] and each value is in the form of (cos(theta_1), sin(theta_1)) pairs up to theta_512. So, 1024 values in total. And of course, x = cos(theta), y = sin(theta).</p>
<p>Then tried to minimize as before:</p>
<ul>
<li>loss_1 = (x ^ 2 + y ^ 2 - 1) ^ 2, to penalize values outside the unit circle</li>
<li>loss_2 = 1 - cos(atan2(y/x) - target), which is 0 if output = target</li>
</ul>
<p>So loss = loss_1 + loss_2 and that’s what I’m back propagating. Target is only a single value for each transducer in radians.</p>
<p>Unfortunately results are the same. My loss never goes below 1.</p>