How to compute autograd for custom pytorch function?

mathematics: <p>I was building custom loss function.</p>
<p>I cant use <code>nn.Module</code> extension because as Ptrblck told, I cant use <code>nn.module</code> when numpy operations is in it,</p>
<aside class="quote no-group" data-username="ptrblck" data-post="51" data-topic="29387">
<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/ptrblck/48/1823_2.png" class="avatar"><a href="https://discuss.pytorch.org/t/custom-loss-functions/29387/51">Custom loss functions</a></div>
<blockquote>
<p>Autograd won’t be able to track the numpy operations, so you would need to implement the backward pass manually via a custom <code>autograd.Function</code></p>
</blockquote>
</aside>
<p>I couldnt use <code>torch.</code> functions because <code>numpy</code> have greater functions where i could try such as <code>np.maximum</code> or things, which is wip in Pytorch. After cleaning ,my final loss function is below</p>
<pre><code class="lang-auto">class TripletLoss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, v1, v2, margin=0.25):
        margin = torch.tensor(margin)
        ctx.save_for_backward(v1, v2, margin)
        scores = np.dot(v1.detach().numpy(), v2.detach().numpy().T)
        batch_size = len(scores)
        positive = np.diag(scores) # the positive ones (duplicates)
        negative_without_positive = scores - 2.0 * np.identity(batch_size)
        closest_negative = negative_without_positive.max(axis=1)
        negative_zero_on_duplicate =  scores * (1.0 - np.eye(batch_size))
        mean_negative = np.sum(negative_zero_on_duplicate, axis=1) / (batch_size - 1)
        triplet_loss1 = np.maximum(0.0, margin - positive + closest_negative)
        triplet_loss2 = np.maximum(0.0, margin - positive + mean_negative)
        triplet_loss = torch.mean(triplet_loss1 + triplet_loss2).requires_grad_()
        return triplet_loss
    
    @staticmethod
    def backward(ctx, grad_output):
        v1, v2, margin = ctx.saved_tensors
        grad_v1, grad_v2, grad_margin = v1, v2, margin
        return grad_v1, grad_v2, grad_margin
</code></pre>
<p>I’m not getting over it , How can i do it when using two or more arguments and numpy functions in it?<br>
It doesnt work when using <code>.backward</code> function,  I tried to use bunch of tactics as described in example such as <code>         grad_input[input[0]&lt;0] = 0</code><br>
some gave error <code>RuntimeError: too many indices for tensor of dimension 0</code>  sometimes on changing codes gave <code>RuntimeError: function TripletLossBackward returned an incorrect number of gradients (expected 2, got 1)</code></p>
<blockquote>
<p><strong>would i get any snippet example for i solving this kind of problem for future cases ,when using numpy functions and tracking using custom autograd functions?</strong></p>
</blockquote>
<p>Also a feature request It would be great if our Pytorch provides <code>numpy</code> equivalent library that contains all numpy  functions as that jax or other frameworks provided.</p>

albanD: <p>Hi,</p>
<blockquote>
<p><code>np.maximum</code> or things, which is wip in Pytorch</p>
</blockquote>
<p>Not sure what you mean by that. You can use torch.max() to get element-wise maximum. Or you can use the threshold function if you just want a threshold as seems to be the case in your code.</p>
<p>I would recommend you read the doc that explains how to write a custom Function here: <a href="https://pytorch.org/docs/stable/notes/extending.html" class="inline-onebox">Extending PyTorch — PyTorch 2.1 documentation</a><br>
That will solve your issues with wrong number and type of returns from the backward.</p>
<blockquote>
<p>Also a feature request It would be great if our Pytorch provides <code>numpy</code> equivalent library that contains all numpy functions as that jax or other frameworks provided.</p>
</blockquote>
<p>This is already work in progress. You can follow the work in this issue: <a href="https://github.com/pytorch/pytorch/issues/38349">https://github.com/pytorch/pytorch/issues/38349</a></p>

mathematics: <p>Final solution for HardTripletloss works pretty well with all <code>numpy</code> eq functions in pytorch , after upgrading Pytorch to unstable version (1.8…) where it consisted all functions i needed to do operation and sticking to <code>nn</code> , found <code>Autograd</code> too hard, will not use it at least for now.</p>
<pre><code class="lang-auto">class HardTripletLoss(nn.Module):
    def __init__(self):
        super(HardTripletLoss, self).__init__()

    def forward(self, v1, v2, margin):
        scores = v1 @ v2.T
        batch_size = len(scores)
        positive = torch.diag(scores)
        negative_without_positive = scores - 2.0 * torch.eye(batch_size)
        closest_negative = negative_without_positive.max(axis=1)[0]
        negative_zero_on_duplicate = scores * (1.0 - torch.eye(batch_size))
        mean_negative = torch.sum(negative_zero_on_duplicate, 1) / (batch_size - 1)
        triplet_loss1 = torch.maximum(margin - positive + mean_negative, torch.tensor(0))
        triplet_loss2 = torch.maximum(margin - positive + closest_negative, torch.tensor(0))
        triplet_loss = torch.mean(triplet_loss2 + triplet_loss1)
        return triplet_loss

</code></pre>
<p>Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>