Batch Hard Triplet Loss: Backpropagation fails due to Function 'SqrtBackward' returned nan values in its 0th output

username: ElisonSherton
cooked: <p>Hi guys!</p>
<p>I have been trying to implement <a href="https://arxiv.org/pdf/1703.07737.pdf" rel="noopener nofollow ugc">this paper</a> which mentions triplet loss with batch hard mining for facial recognition.</p>
<p>Based on my understanding of the paper, I have written the loss function as follows</p>
<pre><code class="lang-auto"># https://arxiv.org/pdf/1703.07737.pdf
import torch
import torch.nn as nn

class batchHardTripletLoss(nn.Module):
    def __init__(self, margin = 0.2, squared = False, agg = "mean"):
        """
        Initialize the loss function with a margin parameter, whether or not to consider
        squared Euclidean distance and how to aggregate the loss in a batch
        """
        super(batchHardTripletLoss, self).__init__()
        self.margin = margin
        self.squared = squared
        self.agg = agg
        self.eps = 1e-8
    
    def get_pairwise_distances(self, feat_vecs):
        """
        Computing distance for every pair using 
        (a - b) ^ 2 = a^2 - 2ab + b^2 
        """
        ab = feat_vecs.mm(feat_vecs.t())
        a_squared = ab.diag().unsqueeze(1)
        b_squared = ab.diag().unsqueeze(0)
        distances = a_squared - 2 * ab + b_squared
        distances.clamp(min = self.eps)
        
        if not self.squared:
            distances = torch.sqrt(distances + self.eps)

        return distances
            
        
    def get_mask(self, labels, type_ = "positive"):
        """
        Get a binary matrix corresponding to valid duplet pairs for 
        (anchor, positive) &amp; (anchor, negative) pairs
        """
        DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
        PK = labels.shape[1]
        mask = torch.zeros(PK, PK).to(DEVICE)
        
        for idx, item in enumerate(labels[0]):
            for inner_idx, inner_item in enumerate(labels[0]):
                
                if type_ == "positive":
                    
                    # Labels should match and the image index shouldn't be the same
                    if (item == inner_item) and (idx != inner_idx):
                        mask[idx, inner_idx] = 1
                elif type_ == "negative":
                    
                    # Labels must be different and image index shouldn't be the same (redundant but still...)
                    if (item != inner_item) and (idx != inner_idx):
                        mask[idx, inner_idx] = 1
                    
        return mask
    
    def forward(self, feat_vecs, labels):
        """
        Define the loss function implementation here
        """
        DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

        # Get the pairwise distances of all images from one another
        distances = self.get_pairwise_distances(feat_vecs)
        
        # Get the toughest positive pair by first filtering out the (anchor, positive)
        # pairs using the get_mask routine and then find the max across rows
        positive_mask = self.get_mask(labels, type_ = "positive").to(DEVICE)
        toughest_positive_distance = (distances * positive_mask).max(dim = 1)[0]
        
        negative_mask = self.get_mask(labels, type_ = "negative").to(DEVICE)
        
        # Add the maxiumum negative distance to all the non-valid pairs
        # on a rowwise basis and then out of them whichever is the minimum 
        # will be our pair distance corresponding to toughest (anchor, negative) pair
        max_negative_dist = distances.max(dim=1,keepdim=True)[0]
        distances = distances + max_negative_dist * (1 - negative_mask).float()
        toughest_negative_distance = distances.min(dim = 1)[0]
        
        # Find the triplet loss by using the two distances obtained above
        triplet_loss = (toughest_positive_distance - toughest_negative_distance + self.margin).clamp(min = self.eps)
        
        # Aggregate the loss to mean/sum based on the initialization of the loss function
        if self.agg == "mean":
            triplet_loss = triplet_loss.mean()
        elif self.agg == "sum":
            triplet_loss = triplet_loss.sum()
            
        return triplet_loss
</code></pre>
<p>However when I go to training, the training runs smoothly for the first few epochs but later, it starts throwing <code>RuntimeError: Function 'SqrtBackward' returned nan values in its 0th output. </code> this error. Here’s a screenshot of the same.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/9/193e83bf30f99db700817933327b6b5a5fe72436.png" data-download-href="https://discuss.pytorch.org/uploads/default/193e83bf30f99db700817933327b6b5a5fe72436" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/193e83bf30f99db700817933327b6b5a5fe72436_2_577x500.png" alt="image" data-base62-sha1="3BjSaHpBvdN86Tdkvxp8IgU3i6i" width="577" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/193e83bf30f99db700817933327b6b5a5fe72436_2_577x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/193e83bf30f99db700817933327b6b5a5fe72436_2_865x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/9/193e83bf30f99db700817933327b6b5a5fe72436.png 2x" data-dominant-color="F4F4F4"><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">898×778 108 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>I tried to <a href="https://discuss.pytorch.org/t/runtimeerror-function-sqrtbackward-returned-nan-values-in-its-0th-output/48702">follow this thread</a> and added  a small epsilon quantity before taking square root to make sure it’s non-zero but it still didn’t work…</p>
<p>Can someone please guide me on how I can tackle this issue? I would be highly obliged.</p>
<p>Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"> !</p>

username: karanjeswani
cooked: <p>The torch clamp function has a different gradient than what you are expecting here. Check this post : <a href="https://stackoverflow.com/questions/60618346/pytorch-autograd-gives-different-gradients-when-using-clamp-instead-of-torch-re" rel="noopener nofollow ugc">Pytorch Autograd gives different gradients when using .clamp instead of torch.relu</a>. Maybe try using relu instead of clamp.</p>

username: ElisonSherton
cooked: <p>Hi <a class="mention" href="/u/karanjeswani">@karanjeswani</a> ,</p>
<p>The issue you pointed to helped!<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/9/890d750d17e78f4d793149727a0205d6737c91b9.png" data-download-href="https://discuss.pytorch.org/uploads/default/890d750d17e78f4d793149727a0205d6737c91b9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/890d750d17e78f4d793149727a0205d6737c91b9_2_690x366.png" alt="image" data-base62-sha1="jyqfdshxmsyNKttkgpmKgTviDaN" width="690" height="366" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/890d750d17e78f4d793149727a0205d6737c91b9_2_690x366.png, https://discuss.pytorch.org/uploads/default/original/3X/8/9/890d750d17e78f4d793149727a0205d6737c91b9.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/9/890d750d17e78f4d793149727a0205d6737c91b9.png 2x" data-dominant-color="EBEBEB"><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">732×389 58.1 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>After substituting torch.clamp with nn.ReLU, the training worked fine as seen above!</p>
<pre><code class="lang-auto"># Find the triplet loss by using the two distances obtained above

# Previously in the loss function:
distances = distances.clamp(min = self.eps)
triplet_loss = (toughest_positive_distance - toughest_negative_distance + self.margin).clamp(min = self.eps)

# Currently in the loss function:
distances = nn.ReLU()(distances)
triplet_loss = nn.ReLU()(toughest_positive_distance - toughest_negative_distance + self.margin)
        
</code></pre>
<p>Thanks a ton!</p>