Edge-classification OOM

[lisiq]

Hi everyone! I am doing edge classification on a big dataset and the GPU is running out of memory even when I'm doing batching.
My guess is that when I am doing the test function, I have to put the data to GPU and that is killing the process.
I have put my code below, any help would be appreciated!

Thank you in advance!

def perform_cross_validation_embedder_heterogeneous(data, parameters):
    n_splits = parameters["n_splits"]
    device = parameters["device"]
    kf = KFold(n_splits=n_splits, random_state=0, shuffle=True)
    print("WARNING: running with a fixed random state")
    output_dict = {}

    for fold, (train_index, test_index) in tqdm(enumerate(kf.split(range(data['student', "takes", "code"].y.size(0))))):

        _, _, test_index, val_index = train_test_split(test_index,
                                        test_index,
                                        test_size=0.5, 
                                        random_state=42)
                                        
        
        train_loader = NeighborLoader(data, 
                                    num_neighbors = {key: [5] for key in data.edge_types}, 
                                    input_nodes=('student', data['student'].node_id),
                                    directed=False,
                                    replace=False,
                                    # drop_last=True,
                                    batch_size=128)
        


        
        model = EmbedderHeterogeneous( 
                n_students =  data["student"].x.size(0),
                n_items = data["code"].x.size(0),
                student_inchannel = data["student"].x.size(1),
                item_inchannel = data["code"].x.size(1),
                hidden_channels=parameters['hidden_dims'],
                edge_channel=data['student', 'takes', 'code'].edge_attr.shape[1],
                metadata=data.metadata(),
                # heads= parameters['heads']
                ).to(device)


        optimizer = torch.optim.Adam(model.parameters(), lr=parameters['learning_rate'], weight_decay=parameters['weight_decay'])

        # Training the model
        losses = []
        best_val_acc = final_test_acc = 0
        early_stopping = 0
        test_val_indices = torch.concat([torch.from_numpy(test_index), torch.from_numpy(val_index)], dim = 0).to(device)

        # data = data.to(device)
        for epoch in tqdm(range(1, parameters['epochs']+1)):
            for batch in tqdm(train_loader): 
                batch = batch.to(device)
                # taking set difference between batch edge ids and union of test and val indexes so that we have no data leakage, i.e., none of the labels in the test and vals are seen by the model, used for the loss etc.
                batch_index = set_diff_1d(batch["student", "takes", "code"].e_id,test_val_indices).to(device)
                loss = train_embedder_heterogeneous(
                    model,
                    batch, #data, batch
                    optimizer,
                    batch_index, #batch_index, train_index
                    )
            losses.append(loss.detach().item())
                
            data = data.to(device)
            train_b = test_embedder_heterogeneous(model, data, train_index, fold, 'train')
            val_b = test_embedder_heterogeneous(model, data, val_index, fold, 'val')
            test_b = test_embedder_heterogeneous(model, data, test_index, fold, 'test')

            if val_b['Balanced Accuracy'+f"_{fold}_val"] > best_val_acc:
                early_stopping = 0

                best_val_acc = val_b['Balanced Accuracy'+f"_{fold}_val"]
                final_test_acc = test_b['Balanced Accuracy'+f"_{fold}_test"]
                best_train_acc = train_b['Balanced Accuracy'+f"_{fold}_train"]

                print(f'\nEpoch: {epoch:03d}, Loss: {loss:.4f}, Train: {best_train_acc:.4f}, '
                    f'Val: {best_val_acc:.4f}, Test: {final_test_acc:.4f}')

                train_b_ = test_embedder_heterogeneous(model, data, train_index, fold, 'train')
                val_b_ = test_embedder_heterogeneous(model, data, val_index, fold, 'val')
                test_b_ = test_embedder_heterogeneous(model, data, test_index, fold, 'test')

                
                # saved_embedding = get_embedding(model, data)
                
            else:
                early_stopping += 1

            if early_stopping == parameters['early_stopping']:
                break
                        
        # Results
        losses_dict = {f'losses_{fold}': losses}
        # embedding = {f'embedding_{fold}': saved_embedding}
        output_dict.update({**parameters,**train_b_, **val_b_, **test_b_,  **losses_dict}) #, **embedding})
        output_dict["device"] = str(output_dict["device"])
    return output_dict


# Train the model function
def train_embedder_heterogeneous(model, data, optimizer, train_index):
    model.train()
    optimizer.zero_grad()
    # from torch import autograd
    # with autograd.detect_anomaly():
    pred = model(
                data=data
                )
    assert pred.isnan().sum() == 0, 'Output'
    target = data['student', 'code'].y[train_index] 
    loss = F.cross_entropy(pred[train_index], target.long())
    loss.backward()
    optimizer.step()
    return loss


# Test the model function
@torch.no_grad()
def test_embedder_heterogeneous(model, data, ids, fold, type):
    model.eval()
    pred = model(
                data=data
                )[ids].cpu()
    target = data['student', 'code'].y[ids].long().cpu().numpy()
    pred = pred.argmax(dim=1, keepdim=True).view(-1).numpy()
    preds = calculate_metrics(target, pred)

    metrics = {k+f"_{fold}_{type}":v for k,v in preds.items()}
    metrics['fold'] = fold
    metrics[f'fold_truths_{fold}_{type}'] = target.tolist()
    metrics[f'fold_preds_{fold}_{type}'] = pred.tolist()
    return metrics

# set difference 
def set_diff_1d(t1, t2, assume_unique=False):
    """
    Set difference of two 1D tensors.
    Returns the unique values in t1 that are not in t2.

    """
    if not assume_unique:
        t1 = torch.unique(t1)
        t2 = torch.unique(t2)
    return (t1[:, None] != t2).all(dim=1)

[lisiq]

This line batch_index = set_diff_1d(batch["student", "takes", "code"].e_id,test_val_indices).to(device) was taking too much memory, hence OOM problem. So i proceeded to split the graph into train, test and val graphs. Now it works as ti should

[rusty1s]

Super :)