Custom loss function (gradient modified by inplace operation)

Allenkei: <p>Hi,</p>
<p>I am trying to optimize a custom loss function for a Graph Neural Network. Given an <strong>n by n</strong> adjacency matrix, I first calculate two network statistics: (1) count of edge, (2) count of 2-stars. Below function returns the 2 network statistics.</p>
<pre><code class="lang-auto">def network_stat(A):
  n = A.shape[0]
  holder = torch.zeros(2, dtype=torch.float32)
  for i in range(n):
    for j in range(n):
      if i &lt; j:
        holder[0] += A[i,j] # first network statistics
        for k in range(n):
          if j &lt; k:
            holder[1] +=A[i,k] * A[j,k] # second network statistics
  return(holder)
</code></pre>
<p>Then my custom loss function for a Graph Neural Network is</p>
<pre><code class="lang-auto">def loss_function(A_pred, eta):
    # minimization: A_pred obtained from Graph Neural Network
    A_pred = (A_pred&gt;0.5).float() # convert to {0,1} adjacency matrix
    predicted_network_stats = network_stat(A_pred)
    loss = torch.dot(eta, sampled_network_stats) # (1 by 2) dot (2 by 1) returns scalar
    return loss
</code></pre>
<p>In other words, I want to learn a Graph Neural Network, whose predicted graph <strong>A_pred</strong> will make the loss function minimized. However, I received below error and I am not sure about how to fix it.</p>
<p><strong>one of the variables needed for gradient computation has been modified by an inplace operation</strong></p>
<p>My understanding is my network_stat(A) function has broken the back-propagation of gradient. Can anyone shed some insight on this problem? I greatly appreciate your help.</p>

ptrblck: <p>Yes, I think your description is correct as <code>network_stat</code> modifies <code>holder</code> inplace in:</p>
<pre><code class="lang-python">holder[0] += A[i,j] # first network statistics
# and
holder[1] +=A[i,k] * A[j,k] # second network statistics
</code></pre>
<p>You could try to replace this operation with its out-of-place version via:</p>
<pre><code class="lang-python">holderA = holderA + A[i, j]
holderB = holderB + (A[i, k] * A[j, k])
...
holder = torch.cat((holderA, holderB), dim=0)
</code></pre>
<p>and check if this would solve the issue.</p>

Allenkei: <p>Hi,</p>
<p>Your solution solves the problem perfectly! Thank you so much again. However, my loss tend to bouncing around. I am wondering if there is anything I did wrong here as I am relatively new to Pytorch. I greatly appreciate any guidance and assistance!</p>
<p>This is my function to calculate 2 graph statistics, which is working fine now:</p>
<pre><code class="lang-auto">def network_stat(adj):
  n = adj.shape[0]
  holderA = torch.tensor([0],dtype=torch.float32)
  holderB = torch.tensor([0],dtype=torch.float32)
  for i in range(n):
    for j in range(n):
      if i &lt; j:
        holderA = holderA + adj[i,j]
        for k in range(n):
          if i &lt; j and j &lt; k:
            holderB = holderB + (adj[i,k] * adj[j,k])
  holder = torch.cat((holderA, holderB), dim=0)
  return(holder)
</code></pre>
<p>This is my function to compute the loss, which depends on the statistics of predicted graph:</p>
<pre><code class="lang-auto">def loss_function(A_pred, eta):

    neg_entropy = torch.log(A_pred).sum()

    A_pred = (A_pred&gt;0.5).float() # make float into 0 and 1 only
    predicted_network_stats = network_stat(A_pred)

    loss = -torch.dot(eta, predicted_network_stats) + neg_entropy
    return loss
</code></pre>
<p>This is my codes to train the model, which involve a minimization (over GCN) within the maximization (over eta):</p>
<pre><code class="lang-auto">for max_epoch in range(20): 
  for min_epoch in range(50):

    optimizer.zero_grad()
  
    A_pred = model(feature, input) 
    loss = loss_function(A_pred, eta)

    loss.backward()
    optimizer.step()  

  with torch.no_grad():
    A_pred = (A_pred&gt;0.5).float()
    expected_network_stats = network_stat(A_pred)
    eta += lr * (observed_network_stats - expected_network_stats)

</code></pre>
<p>Thank you so much!!!</p>

ptrblck: <p>Good to hear it’s working now!<br>
The threshold operation:</p>
<pre><code class="lang-python">A_pred = (A_pred&gt;0.5).float()
</code></pre>
<p>will detach the result from the computation graph so that <code>A_pred</code> and <code>predicted_network_stats</code> will not contribute to the gradient calculation and are just a constant offset to <code>loss</code>.</p>

Allenkei: <p>Thank you very much for pointing out the issue! I have changed the <code>(A_pred&gt;0.5).float()</code> to <code>torch.bernoulli(A_pred)</code> which I think the <code>predicted_network_stats</code> will contribute to the gradient calculation.</p>
<p>By the way, since my function <code>network_stat(adj)</code> involves multiple for-loops iterating over <code>i,j</code> entry of the adjacency matrix, I am wondering if the output from this function will also contribute to the gradient calculation? Or the for-loops also interupt the computation graph for gradient?</p>
<p>Thank you very much again!</p>

ptrblck: <p>Loops do not detach the computation graph and you could check if the output tensor has a valid <code>.grad_fn</code> attribute which indicates that it’s attached to a computation graph.</p>