Directly obtaining gradient w.r.t a node's feature

[YiKangOY]

Hi, I am trying to make a prediction with GNN and obtain gradient w.r.t a node's feature, just like we can do with PyTorch:
import torch
import torch.nn as nn
from torch.autograd import grad
label = 10
lin = nn.Linear(1, 1)
a = torch.Tensor([1]).detach().clone().detach().requires_grad_()
b = lin(a)
c = label-b
print(c.requires_grad) # True
print(a.requires_grad) # True
grad = grad(c, a)
print(grad) # (tensor([-0.8962]),)

But when I want to do it with pytorch geometric, I failed. Here is the example code:
data.x = data.x.clone().detach().requires_grad_()
out= model(data.x, data.edge_index, data.batch)
size = data.x[0][0] # the first feature of the node 0
diff = torch.sum(data.y.to(device) - out)
size_grad = grad(diff, size)[0]

It raises an error says: "RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior."

Here is the GNN I used:
class GCN(torch.nn.Module):
def init(self, input_channels, hidden_channels, layers, dropout_rate):
super(GCN, self).init()
self.node_encoder = Linear(input_channels, hidden_channels)
self.convs = torch.nn.ModuleList()
self.dropouts = torch.nn.ModuleList() # Added dropout layers
for idx in range(layers):
self.convs.append(GCNConv(hidden_channels, hidden_channels))
self.dropouts.append(torch.nn.Dropout(dropout_rate)) # Added dropout layers
self.bns = torch.nn.ModuleList()
for idx in range(layers):
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
self.mlp = torch.nn.Sequential(
torch.nn.Linear(hidden_channels, hidden_channels),
torch.nn.ReLU(),
torch.nn.Linear(hidden_channels, hidden_channels),
torch.nn.ReLU(),
torch.nn.Linear(hidden_channels, 1)
)
def forward(self, x, edge_index, batch):
x = self.node_encoder(x)
for conv, dropout, bn in zip(self.convs, self.dropouts, self.bns): # Added dropout layer
x = conv(x, edge_index)
x = bn(x)
x = x.relu()
x = dropout(x) # Added dropout layer
x = self.mlp(x)
return x

Any help is appreciated, thanks!

[YiKangOY]

Solved, when computing gradient, I guess data.x is as a whole node in the computation graph.
I take size_grad = grad(diff, x)[0], and then take the corresponding index from the gradient tensor, and it works now.