Implementing custom loss from paper

johnnyreimar: <p>Hi all <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
Am trying to implement this paper: <a href="https://arxiv.org/pdf/2003.03522.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/2003.03522.pdf</a> where they employ a detection head to estimate the objects normal distribution, and a regression head to estimate the offset vectors of the objects keypoints. Examplified here: <div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/6/466ff7fa2a4d1d4cfa5d52e0b21e55e5e18b7c99.png" data-download-href="https://discuss.pytorch.org/uploads/default/466ff7fa2a4d1d4cfa5d52e0b21e55e5e18b7c99" title="heads"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/466ff7fa2a4d1d4cfa5d52e0b21e55e5e18b7c99_2_343x500.png" alt="heads" data-base62-sha1="a37jvDkZQvtzO1MVM3mUxEcVHUB" width="343" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/466ff7fa2a4d1d4cfa5d52e0b21e55e5e18b7c99_2_343x500.png, https://discuss.pytorch.org/uploads/default/original/3X/4/6/466ff7fa2a4d1d4cfa5d52e0b21e55e5e18b7c99.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/6/466ff7fa2a4d1d4cfa5d52e0b21e55e5e18b7c99.png 2x" data-dominant-color="B6C3DB"><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">heads</span><span class="informations">421×613 27.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>Have a hard time implementing the regression loss: l1_loss(offset(<strong>p</strong>) + <strong>p</strong> - ground_truth(K_n)) where K_n is the nth keypoint, offset(<strong>p</strong>) is the offset vectors to <strong>p</strong> which is the centroid.</p>
<p>I have the offset vectors and centroid <strong>p</strong>, however during the loss function there are possibly several <strong>p</strong> which need to be added to the offset vectors. So how would I go about this and avoid for-loops?</p>
<p>My current implementation now looks like this:</p>
<pre><code class="lang-auto">class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()

    def forward(self, vector_preds, vector_targets, gauss_preds, gauss_targets):
        vector_targets = vector_targets.permute(0, 2, 3, 1)
        gauss_targets = gauss_targets.permute(0, 2, 3, 1)
        vector_preds = vector_preds.permute(0, 2, 3, 1)
        gauss_preds = gauss_preds.permute(0, 2, 3, 1)  

        E = 0.9
        gauss_mask = gauss_targets &gt; E
        # only want to calculate offsets from centroids (gauss_target &gt; 0.9)
        masked_vector_preds = gauss_mask * vector_preds
        masked_vector_targets = gauss_mask * vector_targets

        vector_loss = F.l1_loss(masked_vector_preds, masked_vector_targets)       

        gauss_loss = F.mse_loss(gauss_preds, gauss_targets)

        loss = vector_loss + gauss_loss
        return loss, vector_loss, gauss_loss
</code></pre>
<p>However, am unsure if i am not just minimizing the vectors now, therefore not taking into account the centroid(s) of the object. Or if, in this case, that even matters?<br>
Hope you can help</p>