Custom Loss Function With Loops

ann81: <p>I am trying to reproduce the loss function from Eq 9 in this <a href="https://arxiv.org/pdf/1812.01217.pdf" rel="noopener nofollow ugc">paper</a>. The idea is that it goes over different matchings of elements in the output and target to create a permutation invariant version of Cross Entropy Loss. This is my implementation currently</p>
<pre><code class="lang-auto">CrossEntropyLoss = nn.CrossEntropyLoss()
def SetCrossEntropyLoss(output, target):
    loss = 0
    for i in range(output.shape[0]):
        se = 0
        for j in range(output.shape[0]):
            H = CrossEntropyLoss(output[i].reshape(1,output.shape[1]), target[j].reshape(1))
            se += torch.exp(-H)
        loss+=torch.log(se)
    return -loss
</code></pre>
<p>I have two questions, though.</p>
<ol>
<li>Would the loops and in place operations mess with the gradients being computed correctly?</li>
<li>Any ideas on how one would make this function more elegant using pytorch operations?</li>
</ol>
<p>Would appreciate any thoughts on either question!</p>

KFrank: <p>Hi Ann!</p>
<aside class="quote no-group" data-username="ann81" data-post="1" data-topic="161852" data-full="true">
<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/a/db5fbb/48.png" class="avatar"> ann81:</div>
<blockquote>
<ol>
<li>Would the loops and in place operations mess with the gradients being computed correctly?</li>
</ol>
</blockquote>
</aside>
<p>Neither the loops nor the inplace operations will cause problems with<br>
backpropagation and gradient computation.</p>
<p>Note, a nice feature of autograd is that it warns you if an inplace<br>
operation has broken backpropagation.  If that were to happen, you<br>
could simply not use inplace operations, e.g.:</p>
<pre data-code-wrap="python"><code class="lang-python">            se = se + torch.exp(-H)
        loss = loss + torch.log(se)
</code></pre>
<p>Even though these looks very similar to your inplace versions, they<br>
create new tensors (at some cost in memory), so no inplace operations<br>
occur.</p>
<aside class="quote no-group">
<blockquote>
<ol start="2">
<li>Any ideas on how one would make this function more elegant using pytorch operations?</li>
</ol>
</blockquote>
</aside>
<p>It is easy enough to implement <code>CrossEntropyLoss</code> yourself, so it is<br>
also easy to implement a loop-free version of <code>SetCrossEntropyLoss</code><br>
using pytorch tensor operations:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; def SetCrossEntropyLoss(output, target):
...     loss = 0
...     for i in range(output.shape[0]):
...         se = 0
...         for j in range(output.shape[0]):
...             H = torch.nn.CrossEntropyLoss() (output[i].reshape(1,output.shape[1]), target[j].reshape(1))
...             se += torch.exp(-H)
...         loss+=torch.log(se)
...     return -loss
...
&gt;&gt;&gt; output = torch.randn (10, 5)
&gt;&gt;&gt; target = torch.randint (5, (10,))
&gt;&gt;&gt;
&gt;&gt;&gt; SetCrossEntropyLoss (output, target)
tensor(-6.9854)
&gt;&gt;&gt;
&gt;&gt;&gt; def SetCrossEntropyLossB (output, target):
...     log_probs = torch.log_softmax (output, dim = 1)
...     xEnts = -log_probs[:, target]
...     loss = -(-xEnts).exp().sum (dim = 1).log().sum (dim = 0)
...     return  loss
...
&gt;&gt;&gt; SetCrossEntropyLossB (output, target)
tensor(-6.9854)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>