Taking different loss based on conditions in a batch

username: legoh
cooked: <p>Hi! I would like to do loss backpropagation through a batch, but this loss value (and hence computation graph) differs per sample, depending on some conditions. I have a way of calculating all my losses, but I’m not sure how to design the backward portion. Let me share a snippet of the code here:</p>
<pre><code class="lang-auto">for _ in range(10000):
    # Mask away first column of all instances to start from there
    mask = torch.zeros((bsz, 20, 20))
    mask_idx = 0
    mask[:, :, mask_idx] = 1.0
    mask = mask.bool()
    # Apply sinkhorn then mask
    # torch.nn.parameter.Parameter is in param_list which I want to optimize via gradient descent
    log_probs = batch_sinkhorn(param_list[0], 0.5, 3, data[1])  
    
    # Find the tour
    # Construct choice mask and combine it with the adjacency matrix mask
    decode_mask = torch.zeros(bsz, nb_nodes)
    decode_mask[:, 0] = 1.0
    decode_mask = decode_mask.bool()
    idx = torch.zeros(bsz).long()
    tour = [torch.zeros(bsz).long()]
    prob_list = []
    
    valid_tour_flag = torch.ones(bsz).bool()
    for _ in range(nb_nodes-1):
        log_prob_choice = log_probs.gather(dim=1, index=idx.unsqueeze(1).unsqueeze(2).repeat(1, 1, nb_nodes)).view(bsz, nb_nodes)

        route_mask = data[1].gather(dim=1, index=idx.unsqueeze(1).unsqueeze(2).repeat(1, 1, nb_nodes)).view(bsz, nb_nodes)
        route_mask = (1.0 - route_mask).bool()

        combined_mask = decode_mask + route_mask # |combined_mask| = (bsz, nb_nodes)
        log_prob_choice = log_prob_choice.masked_fill(combined_mask, -1e29)
        idx = log_prob_choice.argmax(dim=-1)
    
        # Check if a tour is valid by seeing if we chose a masked action
        valid_tour_flag = ~combined_mask.gather(dim=1, index=idx.unsqueeze(1)).squeeze(1) * valid_tour_flag
        valid_tour_flag = valid_tour_flag.bool()
        decode_mask = decode_mask.float().scatter(dim=1, index=idx.unsqueeze(1), src=torch.ones((bsz, 1))).bool()
        tour.append(idx)
        prob_list.append(log_prob_choice.gather(dim=1, index=idx.unsqueeze(1)).squeeze(1))
    tour = torch.stack(tour, dim=1)
    sumLogProb = torch.stack(prob_list, dim=1).sum(-1)
    if torch.any(valid_tour_flag):
            print("There are some valid tours! They are for instances: {}".format(torch.where(valid_tour_flag == True)))
    # Compute tour lengths
    length = compute_tour_length(data[0], tour)
    
    log_probs = log_probs.masked_fill(mask, -1e20)

    # Calculate the dag loss
    d_loss = dag_loss(log_probs)
    optimal_tour_len = data[3]
    
    # This multiplies the log probability of the tour with the tour length and the valid tour flag

    loss = torch.zeros(bsz, requires_grad=True)
    # TODO: find out how to take the loss depending on condition
    loss[torch.where(length &gt; 0)] += length[torch.where(length&gt;0)]
    loss[torch.where(length &lt;= 0)] += dag_loss[torch.where(length&lt;0)]

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
</code></pre>
<p>In general, the loops works as such:</p>
<ul>
<li>I have a bunch of parameters in <code>param_list[0]</code> that I wish to optimize, let’s think of these are just rows of data points with some initial value, with a shape <code>(bsz, 20)</code>
</li>
<li>I apply an algorithm called <code>batch_sinkhorn</code> to get some log probabilities</li>
<li>Based on the log probabilities, I construct a sequence of indexes that I can visit, each time masking out the index already visited</li>
<li>The mask is based on my visits and if the city is even allowed to be visited (captured in <code>route_mask</code>)</li>
<li>After that, I find out the length of all my visits</li>
<li>I calculate my first loss value by using <code>dag_loss</code> which returns a <code>d_loss</code> for each sample</li>
<li>If a sample successfully visited all indexes, I wish to take the loss via <code>sumLogProbs * length</code>
</li>
<li>If a sample is unsuccessful, I wish to take the loss via <code>d_loss</code>
</li>
</ul>
<p>Hopefully this short write-up is clear. Basically, I wish to optimize for one loss if some criterion is met, on samples that were legal. If not, I optimize for another loss for those. Ultimately, all these losses will be pooled together via a mean.</p>

username: ParGG
cooked: <p>Hi, I recently replied to a few similar topics. You can check them out, you should get some insight about how to solve your problem <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"></p>
<aside class="quote quote-modified" data-post="2" data-topic="156279">
  <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/pargg/48/50503_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/conditional-loss-for-multi-class-multi-output-model/156279/2">Conditional loss for multi-class multi-output model</a> 
  </div>
  <blockquote>
    Hi Tal, 
when you load your data, on top of creating the one hot encoded classes, create an additional tensor that will mask your loss. If your room is empty, this tensor will be [1, 0, 0], whilst if it is not empty, it will be like [1, 1, 1]. 
When you compute your final loss (which is probably the sum of loss_r, loss_s, loss_b), just perform the internal (dot) product with the mask I told you about. 
In this case you will have the following situations: 
# empty room, mask = [1, 0, 0]
loss = to…
  </blockquote>
</aside>

<aside class="quote quote-modified" data-post="2" data-topic="155981">
  <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/pargg/48/50503_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/selection-and-application-of-available-methods-based-on-information-from-other-tensor/155981/2">Selection and application of available methods based on information from other tensor</a> 
  </div>
  <blockquote>
    I have a few approaches, but here are two. Let’s start with the one I like the most: 

Loss contribution masking: apply all the functions to the incoming input, and get the resulting vector [func1(input), func2(input), ..., func10(input)] and then perform the internal product (dot product) with the one hot encoded label vector. This will give you as loss just the value of the matching function. Example: class = 2 -&gt; [0, 1, 0, ...], the loss will be loss = 0 * func1(input) + 1 * func2(input) + 0 …
  </blockquote>
</aside>

<p>If not, I can write something more personalised</p>

username: legoh
cooked: <p>I did it! I re-used a flag that I had. Thanks for this nifty idea!</p>