Help with Implementing a custom Loss Function

ayrts: <p>Hello, I am trying to implement this loss function taken from Section 2.1 of <a href="https://arxiv.org/pdf/1703.03717.pdf" rel="noopener nofollow ugc">Right for the Right Reasons: Training Differentiable Models by Constraining their Explanations  (Ross, et al., 2017)</a></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/e/ce8765c68722eb45abf5d4fd73aa425209adda20.png" alt="image" data-base62-sha1="tt2rHGpGxWVYmhaSwGbMyfX3GSI" width="354" height="163"></p>
<p>The first and third term are the Cross-entropy loss and L2 regularization, respectively and are already implemented in Pytorch. The matrix A is a binary mask with dims (Num of samples, W, H, <span class="hashtag-raw">#Color</span> channel). The new loss can be called like this:</p>
<pre data-code-wrap="python"><code class="lang-python">optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=0.01)  # L2 is weight decay

for batch in batches:
   x, target = batch
   logits = model(x)
   loss = explanation_loss(logits, target, l_one, A[index])
   loss.backward()
   optimizer.zero_grad()
   optimizer.step()

</code></pre>
<p>I am confused on how to implement the 2nd term (right reasons).</p>
<p>Specifically:</p>
<ol>
<li>
<p>if my pytorch model outputs logits, I must first convert them to probabilities with softmax and then into a one-hot encoding to take the log, right?</p>
</li>
<li>
<p>How do I implement the input gradient here? (partial derivative w.r.t. x in the image)</p>
</li>
</ol>
<p>Below is my pseudo-code that I want to make functional. <a href="https://github.com/dtak/rrr/blob/master/rrr/multilayer_perceptron.py" rel="noopener nofollow ugc">Here is the original code that I’m having trouble to understand</a> (In the function callled “objective”). Thank you!</p>
<pre data-code-wrap="python"><code class="lang-python">def explanation_loss(logits, target, l_one, mask_matrix):
    ce = nn.CrossEntropyLoss()(logits, target) 
    if l_one==0:
      return ce
    pred = F.softmax(logits, dim=1)
    label_one_hot = torch.nn.functional.one_hot(target, num_classes).float().to(device)
    label_one_hot = torch.clamp(label_one_hot, min=1e-4, max=1.0)  # prevent log(0)

    # missing the input gradient
    right_reasons = l_one* (torch.sum(mask_matrix * torch.log(label_one_hot), dim=1)).mean()
    return ce + right_reasons 
</code></pre>

tom: <ol>
<li>You hardly ever need one hot labels, I don’t think you need them here. Also, I would use torch.nn.functional.cross_entropy_loss instead of re-instantiating the module over and over. You do need predictions, but these are per-class encoded anyways.<br>
You would want to use <code>log_softmax</code> rather than softmax + log.</li>
<li>For the loss involving a gradient w.r.t. the inputs, you’d need to use use <code>requires_grad_()</code> on the inputs and then use <code>gr_input, = torch.autograd.grad(sum_log_hat_y, input) </code> to get the gradient of your loss computation.</li>
<li>If you believe your input gradients to be reasonably independent of the other inputs in the batch (this can be problematic for BatchNorm but is fine for most other things), you can sum over the batch before taking the gradient, this will make things much faster because you only need one gradient.</li>
<li>I think you are missing a **2 before summing.</li>
</ol>
<p>Best regards</p>
<p>Thomas</p>

ayrts: <p>Hello <a class="mention" href="/u/tom">@tom</a>, thank you for the guidance. Now an error is thrown about trying to call the graph a second time when I call loss.backward(). But if I am correct, I just need the gradient values when computing <code>input_grad</code> and the first graph can be dropped.  Do you have further advice on how to solve this without specifying <code>retain_graph=True</code> so I avoid an OOM error?</p>
<p>Below is my current pseudo-code:.</p>
<pre><code class="lang-python">def explanation_loss(model, x, y, l_one):
    logits = model(x)
    ce = F.cross_entropy(logits, target) 
    
    if l_one==0:
        return ce

    sum_log_hat_y = torch.sum(F.log_softmax(logits, dim=1)).mean()
    # same dims as x
    input_grad, = torch.autograd.grad(sum_log_hat_y, x)
    
    mask_matrix = largest_gradient_mask(input_grad)
    mask_matrix = mask_matrix.float()
    
    dot_prod = torch.dot(torch.flatten(mask_matrix), torch.flatten(input_grad))
    
    right_reasons = l_one * (dot_prod) ** 2
    return ce + right_reasons, mask_matrix 


mask_list = []
for batch in clean_train_loader:
    # Mask matrix and loss are calculated in a single pass
    x, target = batch
    x, target = x.to(device).requires_grad_(True), target.to(device)

    optimizer.zero_grad()
    
    loss, A_batch = explanation_loss(model, x, target, 0.01)
    mask_list.append(A_batch)
    
    loss.backward()  # Runtime Error
    optimizer.step()
</code></pre>

tom: <p>In the grad call you would need both retain_graph and create_graph.</p>

ayrts: <p>The loss works as fine as cross entropy when <code>l_one = 0</code>.</p>
<p>However, it does not decrease when I pass a zero matrix. It should behave the same as above because the custom term is multiplied by zero. This is validated by the fact that my assertion statements pass. <a class="mention" href="/u/tom">@tom</a> does <code>autograd.grad</code> influence how the optimizer updates the model parameters, or is this a potential bug for autograd?</p>
<pre><code class="lang-python">
def explanation_loss(logits, x, target, \
        mask_matrix=torch.zeros((len(dataset), 3, 32, 32)), \
        l_one=0.001):

    assert x.requires_grad == True
    ce = F.cross_entropy(logits, target) 
    if l_one == 0:
        return ce

    sum_log_hat_y = torch.sum(F.log_softmax(logits, dim=1)).mean()
    # same dims as x
    input_grad, = torch.autograd.grad(sum_log_hat_y, x, retain_graph=True, create_graph=True)
    
    dot_prod = torch.dot(torch.flatten(mask_matrix), torch.flatten(input_grad))
    assert dot_prod == 0, dot_prod.item()
    
    right_reasons = l_one * ((dot_prod) ** 2)
    assert right_reasons == 0, right_reasons.item()  # These assertions pass
    
    return ce + right_reasons
</code></pre>

tom: <p>I think the torch.dot is wrong. My understanding was that it should be * and then</p>
<pre><code class="lang-auto">right_reasons = l_one * ((dot_prod) ** 2).sum()
</code></pre>
<p>But this doesn’t explain why you would have trouble with the right_reasons term if the mask matrix is all zeros…<br>
Do you have a toy example to demonstrate this, maybe form the paper?</p>

ayrts: <p><a class="mention" href="/u/tom">@tom</a> Oops, you’re right. I was following the original implementation on a perceptron too closely.</p>
<p>It seems that the model updates correctly on vanilla Pytorch, but it doesn’t on Pytorch Lightning. If you work with the latter, I made <a href="https://colab.research.google.com/drive/1lFbFmtQcXlO_QVOOBC7hZWrxuJeJSVee?usp=sharing" rel="noopener nofollow ugc">a notebook</a> for you to reproduce the error, but otherwise I will post this in the PL community. Thanks again!</p>