Segmentation error with SAGE

[bis1999]

So I have been doing my link prediction using https://github.com/snap-stanford/ogb/blob/master/examples/linkproppred/ddi/gnn.py when I am using GCN my code runs perfectly but when I change the use the SAGE code it does not work and give me segmentaion error

/tmp/slurmd/job9389885/slurm_script: line 17: 41836 Segmentation fault (core dumped) python GNN_SAGE_cuda.py

Here is my code as well

import argparse

import faulthandler


faulthandler.enable()
from tqdm import tqdm
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch_geometric.utils import negative_sampling

import torch_geometric.transforms as T

from torch_geometric.nn import GCNConv, SAGEConv

from torch_geometric.nn import SAGEConv

from torch_sparse import SparseTensor
import os 
import pickle
import random
import numpy as np 
from sklearn.metrics import roc_auc_score


device = "cuda" 
#torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)



class GraphSAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(GraphSAGE, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(SAGEConv(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
        self.convs.append(SAGEConv(hidden_channels, out_channels))

        self.dropout = dropout

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, adj_t):
        for conv in self.convs[:-1]:
            x = conv(x, adj_t)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.convs[-1](x, adj_t)
        return x
    
class LinkPredictor(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(LinkPredictor, self).__init__()

        self.lins = torch.nn.ModuleList()
        self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
        self.lins.append(torch.nn.Linear(hidden_channels, out_channels))

        self.dropout = dropout

    def reset_parameters(self):
        for lin in self.lins:
            lin.reset_parameters()

    def forward(self, x_i, x_j):
        x = x_i * x_j
        for lin in self.lins[:-1]:
            x = lin(x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.lins[-1](x)
        return torch.sigmoid(x)

    
def train(model, predictor, x, adj_t,optimizer,negative_ex):


    row, col, _ = adj_t.coo()
    edge_index = torch.stack([col, row], dim=0).to(device)
    edge_index = edge_index.to(device)
    adj_t = adj_t.to(device)
    negative_ex = negative_ex.to(device)


    model.train()
    predictor.train()


    total_loss = total_examples = 0
    
    optimizer.zero_grad()
    
    h = model(x, adj_t)
    

    

    pos_out = predictor(h[edge_index[0]], h[edge_index[1]])
    pos_loss = -torch.log(pos_out + 1e-15).mean()


    neg_out = predictor(h[negative_ex[0]], h[negative_ex[1]])
    neg_loss = -torch.log(1 - neg_out + 1e-15).mean()

    loss = pos_loss + neg_loss
    loss.backward()

    torch.nn.utils.clip_grad_norm_(x, 1.0)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    torch.nn.utils.clip_grad_norm_(predictor.parameters(), 1.0)

    optimizer.step()

   


    return loss.item()


def data_prep(dt_num,idx):
    """Creates the validation,train and test for the model"""
    
    infile_ = open('Link_prediction_proj/Datasets/Data_{}/data_{}.pkl'.format(dt_num,idx), 'rb')
    df = pickle.load(infile_)

    infile = open('Link_prediction_proj/OLP_updated.pickle' ,'rb')
    new= pickle.load(infile)


    
    num_nodes = int(new[new["network_index"] == idx]["number_nodes"])
     
    df_ho_top = df["df_ho_top"]
    df_tr_top = df["df_tr_top"]
    
    # use for model validation (64 links)
    positive_valid= df["Gtr_edges"] 
  
    row_ = [i[0] for i in positive_valid]
    cols_ =[i[1] for i in positive_valid]
    
    #adj training for model validation
    adj_valid = SparseTensor(row=torch.tensor(row_),
                             col=torch.tensor(cols_), sparse_sizes=(num_nodes,num_nodes))
    
    negative_valid_tensor =  torch.tensor(df_tr_top[df_tr_top["Link_present"] ==0]["non_links"].tolist()).T
    postive_valid_tensor = torch.tensor(df_tr_top[df_tr_top["Link_present"] ==1]["non_links"].tolist()).T
    
    all_edges_valid = torch.cat([postive_valid_tensor,negative_valid_tensor],dim =-1)
    all_egdes_label_valid = test_edge_label = torch.cat([
            torch.ones(postive_valid_tensor.size(1)),
            torch.zeros(negative_valid_tensor.size(1))
        ], dim=0)
    
    
    
    # Create the final training data 
    positive_train= df["Gho_edges"] 
    row_train = [i[0] for i in positive_train]
    cols_train =[i[1] for i in positive_train]
    
    #adj training for model validation
    adj_training = SparseTensor(row=torch.tensor(row_train),
                             col=torch.tensor(cols_train), sparse_sizes=(num_nodes,num_nodes))
    
    ## prpared the pytorch htings 
    
    
    test_positive_tensor = torch.tensor(df_ho_top[df_ho_top["Link_present"] ==1]["non_links"].tolist()).T
    test_negative_tensor  =  torch.tensor(df_ho_top[df_ho_top["Link_present"] ==0]["non_links"].tolist()).T
    

    test_edge_label_index = torch.cat(
        [test_positive_tensor,test_negative_tensor],
        dim=-1)
    
    
    test_edge_label = torch.cat([
            torch.ones(test_positive_tensor.size(1)),
            torch.zeros(test_negative_tensor.size(1))
        ], dim=0)
    
    return adj_valid,adj_training,all_edges_valid,all_egdes_label_valid,test_edge_label_index,test_edge_label,negative_valid_tensor,test_negative_tensor
    
    
def get_valid_results(number_l,hid_chan,adj_valid,negative_valid_tensor,all_edges_valid,all_egdes_label_valid):
    
    num_layers = number_l
    hidden_channels = hid_chan
    dropout = 0.5
    lr = 0.005

    device = "cuda"

    print("Entered Valid")








    model = GraphSAGE(hidden_channels, hidden_channels,hidden_channels, num_layers,dropout).to(device)
    emb = torch.nn.Embedding(adj_valid.size(0),hidden_channels).to(device)
    
    predictor = LinkPredictor(hidden_channels, hidden_channels, 1,num_layers, dropout).to(device)
    
    optimizer = torch.optim.Adam(
                list(model.parameters()) +list(predictor.parameters())+list(emb.parameters()), lr=lr)
    

    
    torch.nn.init.xavier_uniform_(emb.weight)
    model.reset_parameters()
    predictor.reset_parameters()

    print("started training")

    for i in range(100):
        ll = train(model, predictor, emb.weight, adj_valid,optimizer,negative_valid_tensor)
       
        
    with torch.no_grad():
        print("no no_grad")
  
        model.eval()
        print("model eval")
        predictor.eval()

        h = model(emb.weight, adj_valid)
        print("some h ")
        pred = predictor(h[all_edges_valid[0]],h[all_edges_valid[1]])

        roc_ = roc_auc_score(all_egdes_label_valid.numpy(),pred.view(-1).detach().cpu().numpy())

    return roc_
    
    
def final_prediction(number_l,hid_chan,adj,neg_ex,all_edges,all_egdes_label):
    num_layers = number_l
    hidden_channels = hid_chan
    dropout = 0.5
    lr = 0.005
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = GraphSAGE(hidden_channels, hidden_channels,hidden_channels, num_layers,dropout).to(device)
    emb = torch.nn.Embedding(adj.size(0),hidden_channels).to(device)
    predictor = LinkPredictor(hidden_channels, hidden_channels, 1,num_layers, dropout).to(device)
    
    optimizer = torch.optim.Adam(
                list(model.parameters()) +list(predictor.parameters())+list(emb.parameters()), lr=lr)
    


    torch.nn.init.xavier_uniform_(emb.weight)
    model.reset_parameters()
    predictor.reset_parameters()
    
    for i in range(100):
        ll = train(model, predictor, emb.weight, adj,optimizer,neg_ex)
        
    with torch.no_grad():
        model.eval()
        predictor.eval()
        h = model(emb.weight, adj)
        pred = predictor(h[all_edges[0]],h[all_edges[1]])
        roc_ = roc_auc_score(all_egdes_label.numpy(),pred.view(-1).detach().cpu().numpy())
        
        
    return roc_,all_egdes_label.numpy(),pred.view(-1).detach().cpu().numpy()



def process_data(dt_n,idx):
    print("data load ")
    adj_valid, adj_training, all_edges_valid, all_egdes_label_valid, test_edge_label_index, test_edge_label, negative_valid_tensor, test_negative_tensor = data_prep(dt_n,idx)
    print("data loaded ")
    num_channels = [2,4,8]
    num_hidden_layers = [32,64,128]
    
    m = {}
    for i in num_channels:
        for j in num_hidden_layers:
            print("something happend")

            m[(i, j)] = get_valid_results(i, j, adj_valid, negative_valid_tensor, all_edges_valid, all_egdes_label_valid)

    pair = max(m, key=m.get)
    hl = pair[0]
    hc = pair[1]

    return final_prediction(hl, hc, adj_training, test_negative_tensor, test_edge_label_index,test_edge_label), m,(hl,hc)


ll = process_data(2,0.0)
print(ll)

'''
with open('Link_prediction_proj/index.pkl', 'rb') as infile:
    index = pickle.load(infile)

# Select a subset of the index for testing (100 in this case)
idx = index





dataset_num = 2
it_ = 1
#int(sys.argv[1])


new_dict = {}
for runs in tqdm(list(range(1,2))):
    print("things started")
    m = {}

    for i in tqdm(idx):
        print(i)
        m[i] = process_data(dataset_num,i)
    new_dict[runs] = m



with open("Link_prediction_proj/Results/SAGE/dat_{}/newfinal20_{}.pkl".format(dataset_num,it_), 'wb')  as file:
    pickle.dump(new_dict , file)
'''

[rusty1s]

Do you know from which line exactly the segmentation fault is coming from?

[bis1999]

It is this line
h = model(emb.weight, adj_valid)

[rusty1s]

I just ran https://github.com/snap-stanford/ogb/blob/master/examples/linkproppred/ddi/gnn.py with the --use_sage option and it works perfectly fine for me. Does it work for you as well?

[bis1999]

Yes this working , so I am not using batchsize , since my datasets are smaller, can that create an issue?

[bis1999]

So, I figured out I was not put my adj_valid on cuda