LinkNeighborLoader error in cuda

[songsong0425]

Hi, always thank you for your effort on maintanance of the package.
I have a simple question about LinkNeighborLoader.

While I run the example code, there was an error about the location of tensors.
I don't know why this error evoked because I think I put all tensors into the device (cuda:0).
Can anyone tell me which tensor is missed?

import torch
import numpy as np
import pandas as pd
from torch_geometric.utils import barabasi_albert_graph
from torch_geometric.loader import DataLoader, LinkNeighborLoader
from torch_geometric.transforms import RandomLinkSplit
from torch_geometric.data import Data

device = "cuda" if torch.cuda.is_available() else "cpu"
np.random.seed(seed=1234)
torch.manual_seed(1234)

torch.cuda.is_available()
# True

# Example graph data generation
bagraph = barabasi_albert_graph(num_nodes=100, num_edges=80)
example_feature = torch.rand(100, 256)
bagraph = bagraph.to(device)
example_feature = example_feature.to(device)

print(bagraph.device)
print(example_feature.device)
# cuda:0
# cuda:0

example_graph = Data(edge_index=bagraph, x=example_feature)
example_graph = example_graph.to(device)
example_graph.is_cuda
# True

transform = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True, split_labels=True)
train_data, val_data, test_data = transform(example_graph)

print(train_data.is_cuda)
print(val_data.is_cuda)
print(test_data.is_cuda)
# True
# True
# True

train_loader = LinkNeighborLoader(train_data, batch_size=16, shuffle=True, neg_sampling_ratio=1.0, num_neighbors=[5,5])
val_loader = LinkNeighborLoader(val_data, batch_size=16, shuffle=True, neg_sampling_ratio=1.0, num_neighbors=[5,5])
test_loader = LinkNeighborLoader(test_data, batch_size=16, shuffle=False, neg_sampling_ratio=1.0, num_neighbors=[5,5])

for batch in loader:
    batch = batch.to(device)
    print(batch)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument tensors in method wrapper_cat)

[songsong0425]

This problem was solved by putting data in the loader into device instead of putting the initial data into cuda from the beginning.

bagraph = barabasi_albert_graph(num_nodes=100, num_edges=80)
example_feature = torch.rand(100, 256)
example_graph = Data(edge_index=bagraph, x=example_feature)

loader = LinkNeighborLoader(example_graph, batch_size=16, shuffle=True, neg_sampling_ratio=1.0, num_neighbors=[5])
for batch in loader:
    batch = batch.to(device)
    print(batch)

# Data(x=[84, 256], edge_index=[2, 221], n_id=[84], e_id=[221], input_id=[16], edge_label_index=[2, 32], edge_label=[32])
# Data(x=[83, 256], edge_index=[2, 207], n_id=[83], e_id=[207], input_id=[16], edge_label_index=[2, 32], edge_label=[32])
# Data(x=[80, 256], edge_index=[2, 215], n_id=[80], e_id=[215], input_id=[16], edge_label_index=[2, 32], edge_label=[32])
# ...

By the way, what is the problem in the original question? I guess that sampled negative edges from LinkNeighborLoader were not in cuda, but it doesn't make sense because main data is in cuda.

[rusty1s]

Let me take a look :)

[rusty1s]

#7572 Should fix this.

[songsong0425]

Thank you for your fast update! I've been trying various things for a few days, and I have a question. In the original question, I split the dataset into train:val:test = 7:1:2 and called them by LinkNeighborLoader. As a result, I got batches as below:

train_loader:
Data(x=[94, 256], edge_index=[2, 382], pos_edge_label=[489], pos_edge_label_index=[2, 489], neg_edge_label=[489], neg_edge_label_index=[2, 489], n_id=[94], e_id=[382], input_id=[32], edge_label_index=[2, 64], edge_label=[64])

val_loader:
Data(x=[95, 256], edge_index=[2, 391], pos_edge_label=[69], pos_edge_label_index=[2, 69], neg_edge_label=[69], neg_edge_label_index=[2, 69], n_id=[95], e_id=[391], input_id=[32], edge_label_index=[2, 64], edge_label=[64])

I thought that from pos_edge_label to neg_edge_label_index were come from RandomLinkSplit, and the rest of them were from LinkNeighborLoader because the number of edge_label_index is double of batch size (I changed the batch size from 16 to 32).

But when I called the data in dataloader by tqdm, the number of samples was the same between train_loader and val_loader (i.e., 31/31). In the below picture, the upper bar is train_loader and the lower bar is val_loader.

I assume that the number of samples between them should be different because they were divided into 7:1.

This is the code for the result. Did I use the wrong variable for the batch iteration? If so, could you please inform me?

# Link prediction task
def train(num_epochs):

    for epoch in range(1, num_epochs+1):

        model.train()
        for data in tqdm(train_loader):
            data = data.to(device)
            optimizer.zero_grad()
        
            z = model(data.x, data.edge_label_index)
            
            out = ((z[data.edge_label_index[0]] * z[data.edge_label_index[1]]).sum(dim=-1)).view(-1)
            tr_loss = criterion(out, data.edge_label)
...
            
        model.eval()
        with torch.no_grad():
            y_val_pred, y_val_true = [], []
            for data in tqdm(val_loader):
                #data = data.to(device)
                y_val_true.append(data.edge_label)
                
                z = model(data.x.to(device), data.edge_label_index.to(device))
                out = ((z[data.edge_label_index[0]] * z[data.edge_label_index[1]]).sum(dim=-1)).view(-1)
                val_loss = criterion(out, data.edge_label.to(device))
...
                
        avg_val_loss = val_losses/len(val_loader.dataset)
        y, pred = torch.cat(y_val_true, dim=0).numpy(), torch.cat(y_val_pred, dim=0).numpy()
        val_f1 = f1_score(y, pred, average='micro')
        
        print(f'Epoch: {epoch:03d}, Training Loss: {avg_tr_loss:.4f}, Validation Loss: {avg_val_loss:.4f}, Validation F1-score: {val_f1:.4f}')

[rusty1s]

It looks like you are missing to pass edge_label_index to the respective LinkNeighborLoader classes.

[songsong0425]

Thank you for your reply! But I'm still confused because I used edge_label_index for out in the dataloader iteration. To clarify the split data(i.e., 489 in the training set and 69 in the validation set), I added some codes as below:

def train(num_epochs):
    best_val_f1 = 0.0

    for epoch in range(1, num_epochs+1):
        tr_losses = 0
        val_losses = 0
        
        model.train()
        for data in tqdm(train_loader):
            data = data.to(device)
            optimizer.zero_grad()
            
            edge_label_index_pn = torch.hstack((data.pos_edge_label_index, data.neg_edge_label_index)) # Added
            edge_label_pn = torch.hstack((data.pos_edge_label, data.neg_edge_label)) # Added
        
            z = model(data.x, edge_label_index_pn)
            
            out = ((z[edge_label_index_pn[0]] * z[edge_label_index_pn[1]]).sum(dim=-1)).view(-1) # Changed to the matched variables
            tr_loss = criterion(out, edge_label_pn) # Changed to the matched variables
            tr_losses += tr_loss.item()
            tr_loss.backward()
            optimizer.step()
        avg_tr_loss = tr_losses/len(train_loader.dataset)
...
# Same modification to the validation step

Unfortunately, it returned the new error RuntimeError: CUDA error: device-side assert triggered. I'm sorry to bother you, but where should edge_label_index is placed to?

[rusty1s]

Can you show me how you initialize the LinkNeighborLoader? You should be doing something similar to

train_loader = LinkNeighborLoader(train_data, edge_label_index=train_data.edge_label_index)

[songsong0425]

Here is the code for LinkNeighborLoader. I realized that I didn't set the edge_label_index=... because I referenced examples/graph_sage_unsup.py code for practice. I greatly appreciate your help. Have a nice day!

transform = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True, split_labels=True)
train_data, val_data, test_data = transform(example_graph)

train_loader = LinkNeighborLoader(train_data, batch_size=32, shuffle=True, neg_sampling_ratio=1.0, num_neighbors=[5,5])
val_loader = LinkNeighborLoader(val_data, batch_size=32, shuffle=True, neg_sampling_ratio=1.0, num_neighbors=[5,5])
test_loader = LinkNeighborLoader(test_data, batch_size=32, shuffle=False, neg_sampling_ratio=1.0, num_neighbors=[5,5])