Apply softmax over multiple segments of output

username: ekarais
cooked: <p>Hi all,</p>
<p>I am faced with the following situation. I am using one model to solve multiple classification tasks, where each classification task itself is multi-class, and the number of possible classes varies across classification tasks. To give an example:</p>
<p>The model outputs a vector with 22 elements, where I would like to apply a softmax over:</p>
<ol>
<li>The first 5 elements</li>
<li>The following 5 elements</li>
<li>The following 8 elements</li>
<li>The last 4 elements</li>
</ol>
<p>This is because the model is simultaneously solving 4 classification tasks, where the first 2 tasks have 5 candidate classes each, the third task has 8 candidate classes, and the final task has 4 candidate classes.</p>
<p>I would also like to define an appropriate cross-entropy loss that follows this same structure.</p>
<p>My questions are:</p>
<ol>
<li>How can I use torch.nn.Softmax to achieve this?</li>
<li>How can I define the custom cross-entropy loss mentioned above?</li>
</ol>
<p>Many thanks!</p>

username: KFrank
cooked: <p>Hello Ege!</p>
<aside class="quote no-group quote-modified" data-username="ekarais" data-post="1" data-topic="95287" 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/e/b77776/48.png" class="avatar"> ekarais:</div>
<blockquote>
<p>I am using one model to solve multiple classification tasks,</p>
<ol>
<li>How can I use torch.nn.Softmax to achieve this?</li>
</ol>
</blockquote>
</aside>
<p>First, for numerical-stability reasons, you shouldn’t use <code>Softmax</code>.<br>
As I outline below, you should use <code>CrossEntropyLoss</code>, which has,<br>
in effect, <code>Softmax</code> built into it.</p>
<aside class="quote no-group">
<blockquote>
<ol start="2">
<li>How can I define the custom cross-entropy loss mentioned above?</li>
</ol>
</blockquote>
</aside>
<p>You don’t need to write a custom cross-entropy loss.  Just use<br>
pytorch’s built-in <code>CrossEntropyLoss</code> four times over, once for<br>
each of your classification tasks.</p>
<p>Your model outputs a batch of prediction vectors of shape <code>[nBatch, 22]</code>.<br>
Your targets could be packaged in a number of ways.  The most<br>
straightforward is probably to have four sets of targets, one for<br>
each classification task.  Let’s call the four tasks <code>A</code>, <code>B</code>, <code>C</code>, and <code>D</code>.<br>
Your targets should be batches of integer class labels.</p>
<p>So, for example, <code>targetA</code> should have shape <code>[nBatch]</code> and consist<br>
of class labels that run from <code>0</code> to <code>4</code>, because task <code>A</code> has five classes.<br>
<code>targetB</code> should be the same.  <code>targetC</code> should also have shape<br>
<code>[nBatch]</code>, but consist of class labels that run from <code>0</code> to <code>7</code> because<br>
task <code>C</code> has eight classes.</p>
<p>Then:</p>
<pre data-code-wrap="python"><code class="lang-python">loss_fn = torch.nn.CrossEntropyLoss()  # only need to do this once

lossA = loss_fn (prediction[:,   0:5], targetA)
lossB = loss_fn (prediction[:,  5:10], targetB)
lossC = loss_fn (prediction[:, 10:18], targetC)
lossD = loss_fn (prediction[:, 18:22], targetD)

loss = lossA + lossB + lossC + lossD
</code></pre>
<p>That is, you use indexing to snip out of your vector of 22 predicted<br>
class labels the set of predictions relevant to each task.  So in this<br>
example, labels 10 through 17 (as indexed by <code>10:18</code>) are the eight<br>
predicted class labels relevant to task <code>C</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: ekarais
cooked: <p>Hello K. Frank,</p>
<p>Thank you for your swift reply. I applied the changes you recommended, and I’m now faced with the following issue:</p>
<p>In the second iteration of the training loop, the following runtime error is produced:</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>
<p>When I specify retain_graph=True when calling backward the first time, and set it to False in the subsequent iterations, I then receive the following error message:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [200, 336]]...
</code></pre>
<p>I have determined that this tensor corresponds to the last layer of my fully connected network. Below, you can find the relevant code snippet from the training loop:</p>
<pre><code class="lang-auto"># update the gradients to zero
        optimizer.zero_grad()

        # forward pass
        logits = model(x)
        
        #compute loss
        for i in range(18):
            loss_array[i] = criterion(logits[:,i*11:(i+1)*11], torch.argmax(labels[:,i*11:(i+1)*11], dim=1))

        for i in range(9):
            loss_array[18+i] = criterion(logits[:,198+i*8:198+(i+1)*8], torch.argmax(labels[:,198+i*8:198+(i+1)*8], dim=1))
        
        for i in range(6):
            loss_array[27+i] = criterion(logits[:,270+i*11:270+(i+1)*11], torch.argmax(labels[:,270+i*11:270+(i+1)*11], dim=1))

        loss = torch.sum(loss_array)

        # backward pass
        loss.backward(retain_graph=first)
        first = False
        train_loss += loss.item()
        
        # update the weights
        optimizer.step()
</code></pre>
<p>The code is a bit more complicated than the example we discussed earlier. I had to use an array to hold all the losses as I have 33 classification tasks instead of 4.</p>
<p>I suspect that I am not using optimizer.step() in the correct way as that is the only operations that updates the network layers.</p>
<p>Would you have an insight into the problem with my code?</p>
<p>Many thanks!<br>
Ege</p>

username: KFrank
cooked: <p>Hello Ege!</p>
<aside class="quote no-group quote-modified" data-username="ekarais" data-post="3" data-topic="95287" 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/e/b77776/48.png" class="avatar"> ekarais:</div>
<blockquote>
<p>In the second iteration of the training loop, the following runtime error is produced:</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>
<pre><code class="lang-auto">        logits = model(x)
        
        #compute loss
        for i in range(18):
            loss_array[i] = criterion(logits[:,i*11:(i+1)*11], torch.argmax(labels[:,i*11:(i+1)*11], dim=1))

        for i in range(9):
            loss_array[18+i] = criterion(logits[:,198+i*8:198+(i+1)*8], torch.argmax(labels[:,198+i*8:198+(i+1)*8], dim=1))
        
        for i in range(6):
            loss_array[27+i] = criterion(logits[:,270+i*11:270+(i+1)*11], torch.argmax(labels[:,270+i*11:270+(i+1)*11], dim=1))

        loss = torch.sum(loss_array)
        
        # update the weights
        optimizer.step()
</code></pre>
</blockquote>
</aside>
<p>I’m not entirely sure what is going on here.  You don’t say what<br>
<code>loss_array</code> is, but since you call <code>torch.sum (loss_array)</code>,<br>
I will assume that <code>loss_array</code> is some kind of pytorch tensor.</p>
<p>If so, indexing into <code>loss_array</code> multiple times could be your<br>
problem.</p>
<p>(You also don’t say what <code>criterion</code> is.  Let me assume that is is<br>
<code>criterion = torch.nn.CrossEntropyLoss()</code>, and therefore that<br>
<code>criterion (...)</code> returns a pytorch tensor of shape <code>[1]</code>, that is<br>
a single number packaged as a tensor.)</p>
<p>Try something like:</p>
<pre data-code-wrap="python"><code class="lang-python">        #compute loss

        loss = 0.0   # python scalar

        for i in range(18):
            # loss will become a pytorch tensor
            loss = loss + criterion(logits[:,i*11:(i+1)*11], torch.argmax(labels[:,i*11:(i+1)*11], dim=1))

        for i in range(18):
             # etc., ...

        loss.backward()

</code></pre>
<aside class="quote no-group">
<blockquote>
<p>When I specify retain_graph=True when calling backward the first time, and set it to False in the subsequent iterations, I then receive the following error message:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [200, 336]]...
</code></pre>
<pre><code class="lang-auto">        # backward pass
        loss.backward(retain_graph=first)
        first = False
        train_loss += loss.item()
</code></pre>
</blockquote>
</aside>
<p>If my theory is right, this use of <code>retain_graph = True</code> is incorrect,<br>
and, rather than fixing the real issue, is just hiding it.  So just call<br>
<code>loss.backward()</code>, as outlined above, without specifying <code>retain_graph</code>.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: Pdieter
cooked: <p>Hi Frank,</p>
<p>thank you for this answer. Would it make sense to weight these loss components? So that lossA would be weighted by 5/22 etc. Even if all tasks (A…D) have the same priority.</p>
<p>Best, Peter</p>