Help Needed with defining Custom Loss Function

user432: <p>Hey, I’ve been recently trying to implement Supervised Contrastive Loss for Semantic Segmentation at the pixel-level. I’m a bit new to PyTorch and I’m finding it difficult to incorporate it to Semantic Segmentation model. Here’s the Supervised Contrastive Loss for Image Classifaction:(<a href="https://arxiv.org/pdf/2004.11362.pdf" rel="nofollow noopener">https://arxiv.org/pdf/2004.11362.pdf</a>)</p>
<pre><code class="lang-auto">class SupConLoss(nn.Module):
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) &lt; 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) &gt; 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss
</code></pre>
<p>My Segmentation model returns output features of shape <strong>[bsz, feat_dim, h, w]</strong> , and the target’s shape and target embedding’s shape are <strong>[bsz, h, w]</strong> &amp; <strong>[bsz, emb_dim, h, w]</strong> respectively.<br>
I’m not sure if this is the right place to post this. Can someone help me with this.</p>
<p>Thanks:)</p>

ptrblck: <p>Do you see any issues, e.g. are errors raised or is the model not converging with this custom loss, or would you just want someone to double check the implementation? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

user432: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, Thanks for the reply. My problem here is SupConLoss, which is defined for Classification problems takes in features of shape 3D, <strong>[bsz, n_views, …]</strong>, and labels of shape 1D, <strong>[bsz]</strong>, but my Segmentation model gives output features of shape 4D, <strong>[bsz, feat_dim, h, w]</strong> and labels are of shape 3D, <strong>[bsz, h, w]</strong>. Is there anyway I can convert this loss so that it fits for my output features shape &amp; labels shape.</p>
<p>Thank you</p>

ptrblck: <p>I’m not sure what the exact difference between the linked paper and <a href="https://arxiv.org/abs/2002.05709">this one</a> is, but the latter was implemented in <a href="https://github.com/HobbitLong/SupContrast">this repository</a>, which might be a good base for you to implement your <code>SupConLoss</code>.<br>
Alternatively, could you point me to the section of the paper, which corresponds to your implementation, so that I could try to figure out the shapes?</p>

user432: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, I guess the only difference between SimCLR &amp; SupContrast is this,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/b/9b8042b7ea80b89082cf4475c165345a8f22b8d0.png" data-download-href="https://discuss.pytorch.org/uploads/default/9b8042b7ea80b89082cf4475c165345a8f22b8d0" title="Screenshot from 2020-07-20 08-22-15"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/b/9b8042b7ea80b89082cf4475c165345a8f22b8d0.png" alt="Screenshot from 2020-07-20 08-22-15" data-base62-sha1="mbCNUwJOK26bJsd584qPksWrqQ8" width="689" height="99" data-dominant-color="DBDBDA"><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 from 2020-07-20 08-22-15</span><span class="informations">1148×165 16 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><br>
I guess Section 3 of the <a href="https://arxiv.org/pdf/2004.11362.pdf" rel="noopener nofollow ugc">SupContrast paper</a> corresponds to the implementation of the loss. Meanwhile, I’m trying to implement the loss into this Semantic Segmentation <a href="https://github.com/zhanghang1989/PyTorch-Encoding" rel="noopener nofollow ugc">repo</a>.</p>
<p>Thank you:)</p>

ptrblck: <p>If the main difference is the number of positive/negative samples, then I would assume the linked code would be reusable. Did you compare it to your code including how the shapes are calculated?</p>

user432: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, The main problem for me is, in the original implementation of SupContrast, <strong>labels</strong> [ground truth] are just numbers from 0-9, but in my case ground truth is an image of shape <strong>[bsz, h,w]</strong>, what I wanted to do is I wanted to define supervised contrastive loss at the pixel level, that is If two pixels are labeled with the same class, I want to make them a positive pair, otherwise, a negative pair.</p>

user432: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a>, I’ve tried this for calculating the pixel-wise loss using the same SupConLoss as above.</p>
<pre><code class="lang-auto">def training(self, epoch):
        train_loss = 0.0
        self.model.train()
        tbar = tqdm(self.trainloader)
        for k, (image, target) in enumerate(tbar):
            self.scheduler(self.optimizer, k, epoch, self.best_pred)
            self.optimizer.zero_grad()
            image, target = image.cuda(), target.cuda()
            features = self.model(image)
            b,c,h,w = features.size()
            for i in range(h):
                for j in range(w):
                    f = features[:,:,i,j]          #[b,c]
                    f = f.unsqueeze(2)             #[b,c,1]
                    t = target[:,i,j]              #[b]
                    loss = self.criterion(f, t)
                    loss.sum().backward(retain_graph=True)
                    self.optimizer.step()
            return loss
            train_loss += loss.item()
            tbar.set_description('Train loss: %.3f' % (train_loss / (k + 1)))
</code></pre>
<p>My GPU utilization is at 95% but I’m not sure if the code is running or not. Can you double check it once. Thanks:)</p>

ptrblck: <p>The code looks alright, but note that <code>train_loss += loss.item()</code> and the print statement won’t be executed, as you are using a <code>return</code> statement before these lines of code.<br>
Besides that I don’t know, if you really need to use <code>retain_graph=True</code>, since you are re-calculating the <code>loss</code> in each iteration.</p>

user432: <p>So I’ve changed my code to this</p>
<pre><code class="lang-auto">def training(self, epoch):
        train_loss = 0.0
        self.model.train()
        tbar = tqdm(self.trainloader)
        for k, (image, target) in enumerate(tbar):
            self.scheduler(self.optimizer, k, epoch, self.best_pred)
            self.optimizer.zero_grad()
            image, target = image.cuda(), target.cuda()
            features = self.model(image)
            b,c,h,w = features.size()
            for i in range(h):
                for j in range(w):
                    f = features[:,:,i,j]
                    f = f.unsqueeze(2)
                    t = target[:,i,j]
                    loss = self.criterion(f, t)
                    loss.sum().backward()
                    self.optimizer.step()
            train_loss += loss.item()
            tbar.set_description('Train loss: %.3f' % (train_loss / (k + 1)))
</code></pre>
<p>And I’ve been getting this error:</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.

</code></pre>

victorc25: <p>Your loop is wrong, you are not changing f and t before running the backward() step again (or at least that’s more or less what that error means).</p>
<p>Comment out the backward() and step() and debug f and t and see what they are doing.</p>

user432: <p>Hey <a class="mention" href="/u/victorc25">@victorc25</a>, Can you help me a bit with how to redefine the loop. I am a bit struck with this. Thanks:)</p>

Hdk: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,<br>
I am working on the same loss function.</p>
<p>But my input to loss function (Zi, Zj, Zk from equation no.4 <a href="https://arxiv.org/pdf/2004.11362.pdf" rel="noopener nofollow ugc">in this paper</a> ) are feature maps which has size of (num_features, h, w) = (128,8,8).</p>
<p>I looked at above code of loss function implementation:<br>
<code>class SupConLoss(nn.Module):</code></p>
<p>But I have no idea what’s going on in that function and moreover their input size is different. So, I was thinking to code the mathematics of that equation which is very complex. I will need several for loops and some logic to differentiate i, j and k.</p>
<p>Batch size=100, so for each i (0 to 99), I have 9 different j values (because in 1 batch there are 9 images having same label as 'i’th image) making k=90(90 images having labels different than that of ‘i’th image). Like wise I have to go through all the images’ feature maps. Can you please suggest how should I proceed?</p>
<p>In addition to that I have doubt on, how to compute dot product of Zi having size of [128,8,8] and Zj of same size which should give me a single value (a scalar value)?</p>

ptrblck: <p>I would recommend to start with creating the right loss function first even if it’s slow in the first version.<br>
I.e. you shouldn’t care too much about nested loop in the first iteration and make sure the outputs and gradients are expected.<br>
Once this is done you can use this slow method as the baseline and try to use vectorized operations to speed it up.</p>

Hdk: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, I have some doubts about the code posted by the person who created this thread.<br>
I could understand till  this line:</p>
<pre><code class="lang-auto"># for numerical stability
</code></pre>
<p>Can you please explain the remaining code?</p>
<p>That loss function (<code>class SupConLoss(nn.Module):</code>) is based on the equation no.4 <a href="https://arxiv.org/pdf/2004.11362.pdf" rel="noopener nofollow ugc">in this paper</a>.</p>
<p>upon comparing the equation and the code, I don’t know where exactly in the code they are <code>separating`` Zi, Zj and Zk</code> values for each Li value and moreover I don’t see any <code>loop</code> which would iterate over all values of i. So, I guess that computes loss only for the <code>first image</code> of the batch.</p>
<p><a href="https://github.com/HobbitLong/SupContrast/blob/master/main_supcon.py" rel="noopener nofollow ugc">In this file</a>, I don’t know what line no.231 (<code>losses.update(loss.item(), bsz</code>) does. I guess it’s just for printing purpose only.</p>
<p>Thank you for sparing your time.</p>

ptrblck: <aside class="quote no-group" data-username="Hdk" data-post="15" data-topic="89524">
<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/h/ec9cab/48.png" class="avatar"> Hdk:</div>
<blockquote>
<p>Can you please explain the remaining code?</p>
</blockquote>
</aside>
<p>I think <a class="mention" href="/u/user432">@user432</a> could better explain the posted code, as I’m unfortunately not familiar with the mentioned paper. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="Hdk" data-post="15" data-topic="89524">
<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/h/ec9cab/48.png" class="avatar"> Hdk:</div>
<blockquote>
<p>I don’t know what line no.231 ( <code>losses.update(loss.item(), bsz</code> ) does. I guess it’s just for printing purpose only.</p>
</blockquote>
</aside>
<p>Yes, the <code>AverageMeter</code> object (which is the class of the <code>losses</code> object) is used for printing purposes and supports calculating the average etc.</p>

avatar: <p>I’m trying to implement this supervised contrastive loss function for the task of semantic segmentation too. Does anybody have any suggestions on where I can find the appropriate implementation of this loss for semantic segmentation? many thanks!</p>