Segmentation fault (core dump) on NeighborLoader

[JiaruiWang]

Hi everyone,

I made a custom dataset.

class USLGCdataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        '''Initialization.
        Args:
            Root is the root data dir './data/'.
        '''
        super().__init__(root, transform, pre_transform, pre_filter)
        self.data = torch.load(osp.join(self.processed_dir, f'us_lgc_in_mem_dataset.pt'))

    @property
    def raw_file_names(self):
        '''Get raw data file names in the './data/raw_dir/'.
        '''

        return ['us_edges_lgc.csv', 
                'us_lgc_2_hop_bfs_voting_label_correct_2nd.csv', 
                'dist_list_idx+52_in_out_all.csv']

    @property
    def processed_file_names(self):
        '''Generated dataset file names saved in the './data/processed_dir/'.
        '''

        return ['us_lgc_in_mem_dataset.pt']


    def download(self):
        # Download to `self.raw_dir`.
        pass
        # path = download_url(url, self.raw_dir)
        # ...

    def process(self):
        idx = 0
        # for raw_path in self.raw_paths:
        #     # Read data from `raw_path`.
        #     # Every sub dir in './data/raw_dir/' for each dataset.

        x = self.get_node_feature('./data/raw_dir/dist_list_idx+52_in_out_all.csv')
        edge_index = self.get_edge_index('./data/raw_dir/us_edges_lgc.csv')
        y = self.get_y_label('./data/raw_dir/us_lgc_2_hop_bfs_voting_label_correct_2nd.csv')

        self.data = Data(x=x, edge_index=edge_index, y=y)
        self.data.num_classes = 52
        self.data.train_mask = torch.Tensor([1 for i in range(5873395)])
        self.data.val_mask = torch.Tensor([1 for i in range(5873395)])
        self.data.test_mask = torch.Tensor([1 for i in range(5873395)])

        torch.save(self.data, osp.join(self.processed_dir, f'us_lgc_in_mem_dataset.pt'))

        
    def get_node_feature(self, file_path: str = None) -> torch.Tensor:
        '''Get node feature matrix with shape [num_nodes, num_node_features].

        Args: 
            file_path: A string of File path for the x_features.csv file.
        Return:
            x: A torch.Tensor with shape (num_nodes, num_node_features).
        '''
        #  Add position anchor sets, node features are the distances to the sets.
        df = pd.read_csv(file_path, sep=',', header=None)
        dfdist = df.iloc[:, 1:157]
        x = torch.Tensor(dfdist.values)
        return x
    
    def get_edge_index(self, file_path: str) -> torch.LongTensor:
        '''Graph connectivity in COO format with shape [2, num_edges].

        Args:
            file_path: A string of File path for the edge_index.csv file.
        Return:
            edge_index: A torch.LongTensor with shape (2, num_edges).
        '''

        df = pd.read_csv(file_path, sep='\t', header=None)
        dft = df.T
        edge_index = torch.LongTensor(dft.values)
        return edge_index

    def get_y_label(self, file_path: str) -> torch.LongTensor:
        '''Node-level ground-truth labels as 243 country classes.

        Args:
            file_path: A string of file path for the c_label.csv contains country labels.
        Return:
            y: A torch.Tensor with shape (num_nodes).
        '''

        df = pd.read_csv(file_path, sep='\t', header=None)
        y = torch.LongTensor(df.T.values[0])
        return y


root = "./data/"
dataset = USLGCdataset(root)
data = dataset.data

Dataset: Data(x=[5873395, 156], edge_index=[2, 84480575], y=[5873395], num_classes=52, train_mask=[5873395], val_mask=[5873395], test_mask=[5873395]):

Number of graphs: 7
Data(x=[5873395, 156], edge_index=[2, 84480575], y=[5873395], num_classes=52, train_mask=[5873395], val_mask=[5873395], test_mask=[5873395])

Number of nodes: 5873395
Number of features: 156
Number of classes: 52
Number of edges: 84480575
Average node degree: 14.38
Number of training nodes: 5873395.0
Number of validation nodes: 5873395.0
Number of testing nodes: 5873395.0
Training node label rate: 1.00

I will get a segmentation fault (core dump) When I run the following lines of code. No other info showed.

loader = torch_geometric.loader.NeighborLoader(
    data,
    # Sample 30 neighbors for each node for 2 iterations
    num_neighbors=[10] * 2,
    # Use a batch size of 128 for sampling training nodes
    batch_size=128,
    input_nodes=data.train_mask,
)

Do you know how it would cause the segmentation fault?
Could it be the network stucture?

[Padarn]

I think this is quite likely to occur if your edge_index refers to nodes that are not in the graph (i.e max(edge_index) > x.shape[0]). Maybe worth checking that as a first step.

[JiaruiWang]

Thank you for your reply. @Padarn
I am using the subgraph edge index from a parent graph. That's why the max(edge_index) > x.shape[0] .
I get the subgraph edge index from the utils.subgraph() function. Do I have to relabel the edge index?
I prefer to keep the original index for the nodes.

[Padarn]

Hi @JiaruiWang - sorry I meant the edge_index in your Data.

        x = self.get_node_feature('./data/raw_dir/dist_list_idx+52_in_out_all.csv')
        edge_index = self.get_edge_index('./data/raw_dir/us_edges_lgc.csv')

If you can provide the dataset I can look in more detail

[JiaruiWang]

Sure.
The original data has 60924683 nodes. The indices are from 0 to 60924682.
The original edge_index (789494545 edges) used node index from 0 to 60924682 as the source and target node index.
The subgraph is a 5873395 nodes subset from the original data. I used the original node indices for the subset nodes. Hence, there are lots of nodes that have indices greater than 5873395 in the subset. Therefore, the edge_index for the subset also has max(edge_index) > 5873395. Is this making sense?
I tried to relabel the edge_index in the subgraph using new nodes indices from 0 to 5873394, without segmentation fault.
Just out of curiosity, is it possible to use the original node indices for the edge_index?

[Padarn]

Got it. You could possibly find a way to do this by setting num_nodes, but then you'll need to adapt any message passing to ignore the nodes that don't actually exist .. it doesn't sound like a very good idea.

Is there any particular reason you want to do this? Maybe there is another way to achieve what you want.

[JiaruiWang]

I see. That was a silly thought. Now I accepted to use the new indices.
Thank you!

[Padarn]

One other way you could achieve this is to set a node_id property on your data before sampling, so that you can later go back to the original ids:

data.node_id = torch.range(0, n_nodes)

Then after sampling you can always do

original_edge_index = torch.cat([
     data.node_id[edge_index[0]],
     data.node_id[edge_index[1]]
])