Need help writing a GNN

[riskiem]

Hello,
I have a set of 35 graphs, with one node feature and one edge feature. They have nodes, that need to be classified into 2 classes.

I am trying to write a GNN of the following form:

Input Layers:
Linear(1,32)
ReLU
Linear(32,32)
ReLU

2 attention layers, using GAT

Output Layer:
Linear(32,32)
ReLU
Linear(32,2)
Softmax

I have so far used the examples here, but have not written a GNN of my own. I am particularly struggling with GAT, since the GATConv doesn't include edge attributes, but I want to include them. Could you please help me write a basic code for this.

Thanks!

[rusty1s]

We have an open pull request for this, see #3078.

Supporting edge features in GATConv is pretty straightforward:

class GATConv(MessagePassing):
    def __init__(self, ...):
        super().__init__(aggr='add')
        self.lin = Linear(in_channels, heads * out_channels)
        self.att = Parameter(torch.Tensor(1, heads, 2 * out_channels + edge_dim)

def forward(self, x, edge_index, edge_attr):
    x = self.lin(x)
    return self.propagate(edge_index, x=x=, edge_attr):

def message(self, x_i, x_j, edge_attr, index, size_i):
    x_i = x_i.view(-1, heads, out_channels)
    x_j = x_j.view(-1, heads, out_channels)
    edge_attr = edge_attr.view(-1, 1, edge_dim)
    out = torch.cat([x_i, x_j, edge_attr], dim=-1)
    alpha = (out * self.att).sum(dim=-1)
    alpha = softmax(alpha, index, num_nodes=size_i)
    return x_j * alpha.unsqueeze(-1)

[rusty1s]

You need to initialize the layer with edge_dim, where edge_dim refers to the feature dimensionality of edge_attr, i.e.: edge_attr.size(-1):

EdgeGATConv(hidden_channels, hidden_channels, edge_dim=edge_dim)

[riskiem]

So I made the following changes:

class Model(torch.nn.Module):
    def __init__(self, hidden_channels):
        super(Model, self).__init__()
        torch.manual_seed(12345)
        edge_dim = edge_attr.size(-1)
        self.lin1 = Linear(data.x.size(-1), hidden_channels)
        self.lin2 = Linear(hidden_channels, hidden_channels)
        self.lin3 = Linear(hidden_channels,2)
        self.conv1 = EdgeGATConv(hidden_channels, hidden_channels, edge_dim = edge_dim)
        self.conv2 = EdgeGATConv(hidden_channels, 2, edge_dim = edge_dim)
        
        
    
    def forward(self, x, edge_index, edge_attr):
    
        
        
        x = self.lin1(x)
        x = x.relu()
        x = self.lin2(x)
        x = x.relu()
        
        x = self.conv1(x, edge_index, edge_attr)
        x = x.relu()
        x = F.dropout(x, p=0.5, training=self.training)
        
        x = self.conv2(x, edge_index, edge_attr)
        
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin2(x)
        x = x.relu()
        x = self.lin3(x)
        
        return x

The train is the same as before:

def train(epoch):
    model.train()
    loss_all = 0
    num_nodes = 0
    for data in train_loader:
        #data = data.to(device)
        optimizer.zero_grad()
        #edge_dim = edge_attr.size(-1)
        out = model(data.x, data.edge_index, data.edge_attr)
        #out = model(data.x, data.edge_index)
        #out = model(data.x)
        loss = criterion(out.squeeze(), data.y.squeeze())
        loss.backward()
        loss_all += loss.item() * data.num_nodes
        num_nodes += data.num_nodes
        optimizer.step()
    #scheduler.step(loss)
    
    return loss_all / num_nodes
    #return num_nodes

def test(loader):
      model.eval()
      test_correct = 0
      pred_len = 0
      for data in loader:
          #edge_dim = edge_attr.size(-1)
          out = model(data.x, data.edge_index, data.edge_attr)
          #out = model(data.x)
          #out = model(data.x, data.edge_index)
          pred = out.argmax(dim= -1)  # Use the class with highest probability.
          test_correct += int((pred == data.y).sum()) # Check against ground-truth labels
          pred_len += pred.numel()
      return test_correct/pred_len

I think the edge_dim problem resolved but there seems to be another problem:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-115-a4704fb6ea19> in <module>()
      1 for epoch in range(1, 250):
----> 2     loss = train(epoch)
      3     tes_acc = test(test_loader)
      4     print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, '
      5           f'Test Accuracy: {tes_acc:.4f}')

4 frames
<ipython-input-113-1f52b9bf1130> in train(epoch)
      7         optimizer.zero_grad()
      8         #edge_dim = edge_attr.size(-1)
----> 9         out = model(data.x, data.edge_index, data.edge_attr)
     10         #out = model(data.x, data.edge_index)
     11         #out = model(data.x)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

<ipython-input-110-263c8ea654c9> in forward(self, x, edge_index, edge_attr)
     21         x = x.relu()
     22 
---> 23         x = self.conv1(x, edge_index, edge_attr)
     24         x = x.relu()
     25         x = F.dropout(x, p=0.5, training=self.training)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

<ipython-input-42-d85d09fc89c9> in forward(self, x, edge_index, size, edge_attr, return_attention_weights)
    136                     fill_or_reduce=self.edge_attr_for_self_loops,
    137                     fill_value=self.edge_attr_fill_value,
--> 138                     num_nodes=num_nodes)
    139             elif isinstance(edge_index, SparseTensor):
    140                 if self.edge_dim is None:

TypeError: add_self_loops() got an unexpected keyword argument 'edge_attr'

[rusty1s]

Keep in mind that you need to call your GNN layers via:

x = self.conv1(x, edge_index, edge_attr=edge_attr)
x = x.relu()
x = F.dropout(x, p=0.5, training=self.training)
x = self.conv2(x, edge_index, edge_attr=edge_attr)

[riskiem]

The error still persists :(

[rusty1s]

Oh I see. Note that the PR also changes other files (like add_self_loops to handle multi-dimensional edge features). You can probably get around with by setting add_self_loops=False in the constructor of your EdgeGATConv.