Torch.autograd.grad got None gradients for cascaded model

username: xuesongwang
cooked: <p>Hi there,<br>
I am trying to build a cascaded model, model1 can be tested in batch (see below).<br>
Then the resulting features of all the samples in the dataset should be concatenated and input to model 2 (kinda like a graph node prediction where all node features generated in model1 have to be aggregated in one graph).</p>
<p>After training, I am expecting to calculate the prediction derivatives in model2 w.r.t the input features in each samples used in model1, and somehow I got stuck and got None gradients. The general pipeline looks similar to the code below, if anyone could give me any idea of what went wrong that would be awesome, thanks <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>
<p>To reproduce:</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch
from torch.utils.data import Dataset, DataLoader

class datasetTest(Dataset):
    def __init__(self):
        super(datasetTest, self).__init__()
        self.x = torch.rand(1024, 3)

    def __getitem__(self, item):
        return self.x[item]

    def __len__(self):
        return self.x.shape[0]

class Model1(nn.Module):
    def __init__(self):
        super(Model1, self).__init__()
        self.l1 = nn.Linear(3, 64)

    def forward(self, x):
        return torch.relu(self.l1(x))

class Model2(nn.Module):
    def __init__(self):
        super(Model2, self).__init__()
        self.l2 = nn.Linear(64, 128)
        self.droput = nn.Dropout()

    def forward(self, x):
        x = torch.relu(self.l2(x))
        x = self.droput(x)
        return x

model1 = Model1().eval() # a pyg model with graph input
model2 = Model2().eval() # model with dropout layers

dataset = datasetTest()
dataloader = DataLoader(dataset, batch_size=16)

feature_list = []
x_list = []
for batch in dataloader:
    batch.requires_grad_(True)
    feature = model1(batch) # model1 on one sample
    x_list.extend(batch)
    feature_list.extend(feature)

feature_dataset = torch.stack(feature_list, dim=0) # get the feature of all the samples in the dataset
feature_dataset.requires_grad_(True)
predict = model2(feature_dataset) # prediction for the whole dataset


# what I want to do: prediction derivative w.r.t. original inputs, got None
deriv = torch.autograd.grad(outputs=predict, inputs=x_list, grad_outputs=torch.ones_like(predict),
                                    allow_unused=True, retain_graph=True, create_graph=True)

# feature_dataset -&gt; original inputs : got None
deriv = torch.autograd.grad(outputs=feature_dataset, inputs=x_list, grad_outputs=torch.ones_like(feature_dataset),
                                    allow_unused=True, retain_graph=True, create_graph=True)


# what I am able to do: predict -&gt; feature_dataset:  has values
deriv = torch.autograd.grad(outputs=predict, inputs=feature_dataset, grad_outputs=torch.ones_like(predict),
                                    allow_unused=True, retain_graph=True, create_graph=True)

</code></pre>

username: InnovArul
cooked: <p>tl: dr;  use <code>append()</code> instead of <code>extend()</code>.</p>
<p><strong>Long story:</strong></p>
<p>The reason for this behavior is due to how <code>extend()</code> works.<br>
<code>extend()</code> iterates over the given <code>batch</code> variable and adds the rows one by one to <code>x_list</code>. By doing this, <code>extend()</code> creates its non-leaf nodes, which are unrelated to the computation graphs involving <code>model1</code>, <code>model2</code>.<br>
Using <code>append()</code> leaves <code>batch</code> variable untouched.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/5/6502df431f460a46eda6e195e3609e5e3a51cfcb.png" alt="image" data-base62-sha1="epAnr1Upuj61GRFZI5ArmUKIRLl" width="456" height="194"></p>
<p>You can find an explanation about leaf variables from <a class="mention" href="/u/alband">@albanD</a> here:</p><aside class="quote quote-modified" data-post="2" data-topic="21751">
  <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/alband/48/215_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/valueerror-cant-optimize-a-non-leaf-tensor/21751/2">ValueError: can't optimize a non-leaf Tensor?</a> 
  </div>
  <blockquote>
    Hi, 
A leaf Variable is a variable that is at the beginning of the graph. That means that no operation tracked by the autograd engine created it. 
This is what you want when you optimize neural networks as it is usually your weights or input. 
So to be able to give weights to the optimizer, they should follow the definition of leaf variable above. 
a = torch.rand(10, requires_grad=True) # a is a leaf variable
a = torch.rand(10, requires_grad=True).double() # a is NOT a leaf variable as it was cr…
  </blockquote>
</aside>


username: xuesongwang
cooked: <p>Thank you very much for the detailed explanation !!<br>
It worked when I changed <code>x_list.extend(batch)</code> into <code>x_list.append(batch)</code>  and write the derivative like this:</p>
<pre><code class="lang-auto">deriv = torch.autograd.grad(outputs=predict, inputs=x_list, grad_outputs=torch.ones_like(predict),
                                    allow_unused=True, retain_graph=True, create_graph=True)
</code></pre>
<p>the result is a tuple as well (length equals <code>len(x_list)</code>), I can cat them together to get the final derivative.</p>
<p>However, if I did this then it didn’t work, I guess this cat function create a new variable that is unrelated to model1 again?</p>
<pre><code class="lang-auto">x_data = torch.cat(x_list, dim=0)
feature_dataset = torch.stack(feature_list, dim=0) # get the feature of all the samples in the dataset
feature_dataset.requires_grad_(True)
predict = model2(feature_dataset) # prediction for the whole dataset
deriv = torch.autograd.grad(outputs=predict, inputs=x_data, grad_outputs=torch.ones_like(predict),
                                    allow_unused=True, retain_graph=True, create_graph=True)
</code></pre>

username: InnovArul
cooked: <p>yes that’s right. <code>cat()</code> creates a new non-leaf node as well.</p>