Testing samplers on node classification task

[MichaelHopwood]

Was hoping to get some feedback as I am new to this library (awesome work btw!!).
I adapted the benchmark/citation/ code to test out different training techniques on a node classification task.

I added some code which defines the train_loader and utilizes it for batching the training.

def run_train_batched(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
              sampler, num_neighbors, batch_size,
              permute_masks=None, logger=None):
    
    val_losses, accs, durations = [], [], []
    for _ in range(runs):
        data = dataset[0]
        
        if permute_masks is not None:
            data = permute_masks(data, dataset.num_classes)
        
        train_data = data.subgraph(data.train_mask)
        
        if sampler == 'LinkNeighborLoader':
            train_loader = LinkNeighborLoader(
                train_data,
                num_neighbors=num_neighbors,
                batch_size=batch_size,
                edge_label_index=train_data.edge_index,
            )
        elif sampler == 'NeighborLoader':
            train_loader = NeighborLoader(
                train_data,
                num_neighbors=num_neighbors,
                batch_size=batch_size,
            )
        else:
            raise Exception(".")

        data = data.to(device)

        model.to(device).reset_parameters()
        optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        t_start = time.perf_counter()

        best_val_loss = float('inf')
        test_acc = 0
        val_loss_history = []

        for epoch in range(1, epochs + 1):
            
            for batch in train_loader:
                train(model, optimizer, batch)
            
            eval_info = evaluate(model, data)
            eval_info['epoch'] = epoch
            
            if logger is not None:
                logger(eval_info)

            if eval_info['val_loss'] < best_val_loss:
                best_val_loss = eval_info['val_loss']
                test_acc = eval_info['test_acc']

            val_loss_history.append(eval_info['val_loss'])
            if early_stopping > 0 and epoch > epochs // 2:
                tmp = tensor(val_loss_history[-(early_stopping + 1):-1])
                if eval_info['val_loss'] > tmp.mean().item():
                    break

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        t_end = time.perf_counter()

        val_losses.append(best_val_loss)
        accs.append(test_acc)
        durations.append(t_end - t_start)
    loss, acc, duration = tensor(val_losses), tensor(accs), tensor(durations)

    print(f'Val Loss: {float(loss.mean()):.4f}, '
          f'Test Accuracy: {float(acc.mean()):.3f} ± {float(acc.std()):.3f}, '
          f'Duration: {float(duration.mean()):.3f}s')

I then look at the results when using these different samplers, as well as comparing it to training without batching (the default benchmark/citation/ run_train method).
I'm seeing poor results when using LinkNeighborLoader so I assume there's a gap in my knowledge or something wrong with my implementation.

+ python gcn.py --dataset cora --runs 20 --random_splits
Val Loss: 0.8895, Test Accuracy: 0.797 ± 0.018, Duration: 3.597s

+ echo 'If NeighborLoader, what neighbor config?'
If NeighborLoader, what neighbor config?
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 5 10 --batch_size 128
Val Loss: 1.3087, Test Accuracy: 0.718 ± 0.040, Duration: 1.560s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 5 10 25 --batch_size 128
Val Loss: 1.2759, Test Accuracy: 0.728 ± 0.039, Duration: 1.959s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 20 50 --batch_size 128
Val Loss: 1.2742, Test Accuracy: 0.724 ± 0.039, Duration: 1.632s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 --batch_size 128
Val Loss: 1.2777, Test Accuracy: 0.730 ± 0.029, Duration: 1.701s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 50 --batch_size 128
Val Loss: 1.2684, Test Accuracy: 0.737 ± 0.032, Duration: 1.772s

+ echo 'If LinkNeighborLoader, what neighbor config?'
If LinkNeighborLoader, what neighbor config?
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 5 10 --batch_size 128
Val Loss: 1.9173, Test Accuracy: 0.145 ± 0.099, Duration: 1.072s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 5 10 25 --batch_size 128
Val Loss: 1.9236, Test Accuracy: 0.138 ± 0.087, Duration: 1.127s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 20 50 --batch_size 128
Val Loss: 1.9423, Test Accuracy: 0.120 ± 0.044, Duration: 1.057s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 --batch_size 128
Val Loss: 1.9330, Test Accuracy: 0.142 ± 0.085, Duration: 1.180s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 50 --batch_size 128
Val Loss: 1.9449, Test Accuracy: 0.120 ± 0.070, Duration: 1.093s

+ echo 'If NeighborLoader, what batch size?'
If NeighborLoader, what batch size?
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 --batch_size 16
Val Loss: 1.1477, Test Accuracy: 0.738 ± 0.025, Duration: 3.039s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 --batch_size 32
Val Loss: 1.1948, Test Accuracy: 0.721 ± 0.039, Duration: 2.248s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 --batch_size 64
Val Loss: 1.2167, Test Accuracy: 0.727 ± 0.033, Duration: 1.905s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 --batch_size 128
Val Loss: 1.2731, Test Accuracy: 0.729 ± 0.043, Duration: 1.870s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler NeighborLoader --num_neighbors 50 50 --batch_size 256
Val Loss: 1.2526, Test Accuracy: 0.719 ± 0.022, Duration: 1.939s

+ echo 'If LinkNeighborLoader, what batch size?'
If LinkNeighborLoader, what batch size?
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 --batch_size 16
Val Loss: 1.9570, Test Accuracy: 0.094 ± 0.028, Duration: 1.384s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 --batch_size 32
Val Loss: 1.9350, Test Accuracy: 0.140 ± 0.088, Duration: 1.282s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 --batch_size 64
Val Loss: 1.9414, Test Accuracy: 0.125 ± 0.060, Duration: 1.069s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 --batch_size 128
Val Loss: 1.9268, Test Accuracy: 0.149 ± 0.097, Duration: 1.097s
+ python gcn.py --dataset cora --runs 20 --random_splits --sampler LinkNeighborLoader --num_neighbors 50 50 --batch_size 256
Val Loss: 1.9289, Test Accuracy: 0.153 ± 0.086, Duration: 1.132s

[rusty1s]

The LinkNeighborLoader is only applicable for edge-level tasks, e.g., link prediction. Since you are working on a node-level task, it isn't suitable in this scenario.

For the NeighborLoader, note that you will need to perform slicing within train() such that the loss is only computed for the seed nodes you are sampling neighbors for. See here for an example.