Model not learning with custom loss function and parameters not updated

JFM_AI: <p>Dear community,</p>
<p>I have an ML model (A GNN to be more precise) with a softmax activation function in its output layer, with num_classes = K.</p>
<p>They way my custom loss function works is as follows (summary):</p>
<ul>
<li>Takes the output of the whole dataset.</li>
<li>Samples a class for each data point.</li>
<li>Build a graph structure with max K nodes.</li>
<li>Runs an evaluation algorithm of my own that return a Value. The higger the value, the better the classification.</li>
</ul>
<p>Then, I use Pytorch loss.backward() to update the gradients.</p>
<pre><code class="lang-auto">def train(data):
    optimizer.zero_grad()  # Clear gradients.
    out = model(data.x, data.edge_index)  # Perform a single forward pass.
    L = sampling_classes(values, torch.exp(out))
    graph = build_graph(L)
    loss = loss_value_function(graph, out)
    loss.backward()  # Derive gradients.
    optimizer.step()  # Update parameters based on gradients.
    return loss
</code></pre>
<pre><code class="lang-auto">def loss_value_function(graph, out):
    val = evaluate_graph(graph)
    return (torch.sum(torch.exp(out)) - torch.sum(torch.tensor(val)))
</code></pre>
<p>In <code>return (torch.sum(torch.exp(out)) - torch.sum(torch.tensor(val)))</code>, I realised I have to substract <code>torch.sum(torch.tensor(val))</code> to <code>(torch.sum(torch.exp(out))</code> (which is always the number of data points), otherwise I break the gradient, and then the parameters are not updated if I only return <code>(-1) * torch.sum(torch.tensor(val)))</code>. Is this correct? is there any way I can just return <code>(-1) * torch.sum(torch.tensor(val)))</code> and then get the gradient updated when I perform the backward pass?</p>
<p>I tried returng <code>(-1) * torch.sum(torch.tensor(val)))</code> as loss and then doing <code>loss = torch.tensor(loss, required_grad=True)</code>, but the parameters are never updated.</p>
<p>Any suggestion would be appreaciated.</p>
<p>Thank you!</p>
<p>Cheers.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/jfm_ai">@JFM_AI</a>,</p>
<aside class="quote no-group" data-username="JFM_AI" data-post="1" data-topic="163231">
<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/j/a183cd/48.png" class="avatar"> JFM_AI:</div>
<blockquote>
<p>(torch.tensor(val)</p>
</blockquote>
</aside>
<p>When you do <code>torch.tensor(val)</code> you’re breaking your computational graph, which is why your parameters are not updating. So, remove all these <code>torch.tensor</code> and re-run your code and see if it updates then!</p>

JFM_AI: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>, Thank you for your reply. I have tried this, however, if I remove <code>torch.tensor()</code>, then when I do <code>loss.backward()</code> I get this <code>Error: AttributeError: 'numpy.float64' object has no attribute 'backward'</code>:</p>
<pre><code class="lang-auto">def loss_value_function(graph, out):
    val = evaluate_graph(graph)
    return np.sum(val)
</code></pre>

ptrblck: <p>You cannot use numpy arrays directly as this will also break the computation graph. Either use PyTorch operations only or write a custom <code>autograd.Function</code> if you need to use numpy operations.</p>

JFM_AI: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> , thank you for your response. That’s true. I also tried <code>return torch.sum(torch.tensor(VK, requires_grad=True))</code>, but the computational graph breaks and the parameters don’t update as well.<br>
However, if I do <code>return (torch.sum(torch.exp(out)) - torch.sum(torch.tensor(VK)))</code>, then the parameters are updated. Although I find this pretty weird. Any suggestion?</p>

ptrblck: <p>Recreating a tensor will break the computation graph as already mentioned. If your second approach works fine I would guess that another tensor might be attached to a computation graph (<code>out</code>) while <code>VK</code> acts as a constant and won’t affect the gradient calculation.</p>

JFM_AI: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> .</p>
<p>Yes, this is what I thought. out is a tensor, and it is constant. Then, I just use VK as a value to be substracted from out. The greate VK, the smaller the loss.</p>
<p>Thank you</p>

JFM_AI: <p>Ok, I found what’s the problem. The problem is that I was doing a random sampling in my log_softmax output layer and that, breaks the gradient. In order to update the parameters when we want to do a  random sampling, we need to create a categorical distribution, sample from that distribution and then, update the parameters with the log derivative:</p>
<pre><code class="lang-auto">def train(data):
    optimizer.zero_grad()  # Clear gradients.
    out = model(data.x, data.edge_index)  # Perform a single forward pass.
    m = torch.distributions.categorical.Categorical(out) # Creates a categorical distribution
    L = m.sample() # Samples from the distribution. Each datapoint is assigned to a class.
    graph = build_graph(L)
    val = loss_value_function(graph)
    loss  = -m.log_prob(L) * val 
    loss.sum().backward()  # Derive gradients.
    optimizer.step()  # Update parameters based on gradients.
    return loss
</code></pre>
<pre><code class="lang-auto">def loss_value_function(graph):
    val = evaluate_graph(graph)
    return val
</code></pre>
<p>Hope this helps.</p>
<p>Cheers,</p>