Link Neighbor Loader explanation

[sciencewu]

Hi, huge fan of PyG. Also of the discussions here, very helpful!
I tried to join the Slack channel but it says there is an error.
My question is if I do this: train_data = Batch.from_data_list(dataset[train_index]) and then LinkNeighborLoader will the negative edges be just within the graphs or with the other graphs?

I have added my code just to check what I might be doing wrong as the predictions are not as good:

# GNN model
class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, heads, dropout, edge_dim):
        super().__init__()
        self.dropout = dropout

        self.conv1 = GATConv(in_channels = in_channels, 
                            out_channels = hidden_channels, 
                            heads = heads, 
                            edge_dim = edge_dim,
                            dropout=dropout
                            )
        self.conv2 = GATConv(in_channels = hidden_channels * heads, 
                            out_channels = out_channels,
                            edge_dim = edge_dim,
                            dropout=dropout
                            )

    def forward(self, x, edge_index, edge_attr):
        # x = F.dropout(x, p=self.dropout, training=self.training)
        x = F.elu(self.conv1(x, edge_index, edge_attr))
        # x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv2(x, edge_index, edge_attr)
        return x

# Train-Test functions
def train(model, loader, optimizer):
    model.train()

    total_loss = 0
    for data in loader:
        data = data.to(device)
        optimizer.zero_grad()
        loss = F.cross_entropy(model(data.x, data.edge_index, data.edge_attr), data.y.long())
        total_loss += loss.item() * data.num_graphs
        loss.backward()
        optimizer.step()
    return total_loss / len(loader.dataset)

@torch.no_grad()
def encode(model, loader):
    model.eval()

    xs, ys = [], []
    for data in loader:
        data = data.to(device)
        out = model(data.x, data.edge_index, data.edge_attr)
        xs.append(out.argmax(dim=1))
        ys.append(data.y.cpu())
    return torch.cat(xs, dim=0).long().detach().cpu(), torch.cat(ys, dim=0).long().cpu()


# Test function
@torch.no_grad()
def test(model, train_loader, test_loader, fold):
    model.eval()

    x, y = encode(model, train_loader)
    pred = calculate_metrics(x.cpu(), y.cpu())

    metrics = {k+f"_{fold}_train":v for k,v in pred.items()}
    metrics['fold'] = fold
    metrics[f'fold_truths_{fold}_train'] = y.detach().cpu().numpy().tolist()
    metrics[f'fold_preds_{fold}_train'] = x.detach().cpu().numpy().tolist()

    x, y = encode(model, test_loader)
    
    pred = calculate_metrics(x.cpu(), y.cpu())

    metrics.update({k+f"_{fold}_test":v for k,v in pred.items()})
    metrics[f'fold_truths_{fold}_test'] = y.detach().cpu().numpy().tolist()
    metrics[f'fold_preds_{fold}_test'] = x.detach().cpu().numpy().tolist()

    return metrics

# Training procedure
kf = KFold(n_splits=5, random_state=0, shuffle=True)
output_dict = {}

for fold, (train_index, test_index) in tqdm(enumerate(kf.split(dataset))):

    # Group all training graphs into a single graph to perform sampling:
    train_data = Batch.from_data_list(dataset[train_index])
    sampler = ImbalancedSampler(train_data.y.long())
    loader = NeighborLoader(
                            train_data, 
                            batch_size=1024, 
                            #shuffle=True,
                            num_neighbors=[10, 10],
                            num_workers=1, 
                            persistent_workers=True,
                            sampler=sampler
                            )
    
    test_loader = DataLoader(dataset[test_index], batch_size=parameters['batch_size'])
    train_loader = DataLoader(dataset[train_index], batch_size=parameters['batch_size'])
    
    model = GAT(
        in_channels= dataset.num_features,
        out_channels=2,
        hidden_channels=parameters['hidden_dims'],
        edge_dim=dataset.num_edge_features,
        dropout=parameters['dropout'],
        heads=parameters['heads']
    ).to(device)
    optimizer = torch.optim.Adam(
                                 model.parameters(), 
                                 lr=parameters['learning_rate'], 
                                 weight_decay=parameters['weight_decay']
                                 )
    
    
    losses = []
    for epoch in tqdm(range(1, parameters['epochs'] + 1)):
        loss = train(model, loader, optimizer)
        losses.append(loss)

    
    # Results
    eval_ = test(model, train_loader, test_loader, fold)
    losses_dict = {f'losses_{fold}': losses}
    output_dict.update({**parameters,**eval_,  **losses_dict})

[rusty1s]

What problems do you have when trying to join the Slack?

LinkNeighborLoader acts like working on a single graph, so it will also create negatives between different graphs in that way. How big are your graphs though? Do you desperately need to make use of LinkNeighborLoader or can't you just iterate over your graphs in mini-batches (in case they are small).

[sciencewu]

Hi! I am sorry for the delay on the answer I had to reproduce the model from the scratch.
So the code that I am using is in the updated question now.
My concern is why is the model not performing as good as only using the node features only (random forest)?
I am trying to predict a class for each node in the graph, but I'm training on several graphs.

[rusty1s]

Does the same hold if you use SAGEConv rather than GATConv?

[lisiq]

SAGEConv is performing better than GATConv here, it depends on the run but usually outperforms by 5-6%.
Which is still ~20% worse than only using the node features as tabular data with random forest 🤔

[rusty1s]

Mh, that is weird. SAGEConv has a built-in skip-connection, which means that it theoretically could learn to not use any neighbors (resembling an MLP). It seems to struggle to do that, so it actually looks like your graph structure hurts the underlying model performance.

[lisiq]

Yes, that's what I was thinking as well. I will get a new batch of data which may be a bit different and if you are interested I will keep you noted:)

But DeepGCN is performing way better, only ~2 % worse than random forest.