Node level prediction with heterogenous GCN

[PietroLodiRizzini]

I'm trying to train a GNN to predict the price of apartments by taking into account the features of the neighbourhood.
I have started with this simple graph structure built from tabular data:

HeteroData(
  listing={
    x=[8528, 27],
    y=[8528],
    train_mask=[8528],
    val_mask=[8528],
    test_mask=[8528],
  },
  poi={ x=[3157, 7] },
  (listing, is_near, poi)={
    edge_index=[2, 104254],
    edge_weight=[104254],
  },
  (poi, rev_is_near, listing)={
    edge_index=[2, 104254],
    edge_weight=[104254],
  }
)

where i have 2 node types, listing and poi. The task is to predict data['listing'].y which is the price of the listing.

This is what I have so far:

class HeteroGNN(torch.nn.Module):
  def __init__(self, metadata, hidden_channels, out_channels, num_layers):
    super().__init__()
    self.convs = torch.nn.ModuleList()

    for _ in range(num_layers):
        conv = HeteroConv({edge_type: GraphConv((-1, -1), hidden_channels) for edge_type in data.edge_types
                          }, aggr='sum')
        self.convs.append(conv)

    self.lin = Linear(hidden_channels, out_channels)

  def forward(self, x_dict, edge_index_dict):
      for conv in self.convs:
        x_dict = conv(x_dict, edge_index_dict, edge_weight_dict=data.edge_weight_dict)
        x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()}


      return self.lin(x_dict['listing'])

model = HeteroGNN(data.metadata(), hidden_channels=64, out_channels=1, num_layers=2)
optimizer = torch.optim.Adam(model.parameters(), lr=.0001, weight_decay=0.001)

def train():
  model.train()
  optimizer.zero_grad()
  out = model(data.x_dict, data.edge_index_dict)
  mask = data['listing'].train_mask
  target = data['listing'].y[mask].unsqueeze(1)
  #print("out=", out[mask], "target=", target)
  
  loss = F.mse_loss(out[mask], target) 
  loss.backward()
  optimizer.step()
  return float(loss)

@torch.no_grad()
def test():
  model.eval()
  pred = model(data.x_dict, data.edge_index_dict).argmax(dim=-1)

  accs = []
  for split in ['train_mask', 'val_mask', 'test_mask']:
      mask = data['listing'][split]
      print("pred", pred[mask][:10], "actual", data['listing'].y[mask][:10])
      acc = sklearn.metrics.mean_absolute_error(data['listing'].y[mask], pred[mask])
      accs.append(float(acc))
  return accs


for epoch in range(1, 101):
  loss = train()
  train_acc, val_acc, test_acc = test()
  print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, ' f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')

The problem is that the predicted values are always 0 (i.e. [0, 0, ..., 0])
Any suggestion or advice would be strongly appreciated

[rusty1s]

Since this is a regression use-case,

pred = model(data.x_dict, data.edge_index_dict).argmax(dim=-1)

looks incorrect to me? I think you just want to drop the argmax here.

[PietroLodiRizzini]

That solved the problem of null outputs. Although, the model is performing very poorly.
I wanted to try to compute listing node embeddings with GraphSage to increase the features of the listings dataset and eventually train a regression model over the augmented dataset. This should increase the performance compared to keeping the original features only. Would this be feaseble on this structure?
Any other suggestion is strongly appreciated

[rusty1s]

Yes, that should work. However, if your GraphSAGE model underperforms against a simple MLP, there might be other issues going on.