Source code for torch_geometric.datasets.linkx_dataset

import os.path as osp
from typing import Callable, List, Optional

import numpy as np
import torch
import torch.nn.functional as F

from import Data, InMemoryDataset, download_url

[docs]class LINKXDataset(InMemoryDataset): r"""A variety of non-homophilous graph datasets from the `"Large Scale Learning on Non-Homophilous Graphs: New Benchmarks and Strong Simple Methods" <>`_ paper. .. note:: Some of the datasets provided in :class:`LINKXDataset` are from other sources, but have been updated with new features and/or labels. Args: root (string): Root directory where the dataset should be saved. name (string): The name of the dataset (:obj:`"penn94"`, :obj:`"reed98"`, :obj:`"amherst41"`, :obj:`"cornell5"`, :obj:`"johnshopkins55"`, :obj:`"genius"`). transform (callable, optional): A function/transform that takes in an :obj:`` object and returns a transformed version. The data object will be transformed before every access. (default: :obj:`None`) pre_transform (callable, optional): A function/transform that takes in an :obj:`` object and returns a transformed version. The data object will be transformed before being saved to disk. (default: :obj:`None`) """ github_url = ('' 'raw/master/data') gdrive_url = '' facebook_datasets = [ 'penn94', 'reed98', 'amherst41', 'cornell5', 'johnshopkins55' ] datasets = { 'penn94': { 'data.mat': f'{github_url}/facebook100/Penn94.mat' }, 'reed98': { 'data.mat': f'{github_url}/facebook100/Reed98.mat' }, 'amherst41': { 'data.mat': f'{github_url}/facebook100/Amherst41.mat', }, 'cornell5': { 'data.mat': f'{github_url}/facebook100/Cornell5.mat' }, 'johnshopkins55': { 'data.mat': f'{github_url}/facebook100/Johns%20Hopkins55.mat' }, 'genius': { 'data.mat': f'{github_url}/genius.mat' }, 'wiki': { '': f'{gdrive_url}id=1p5DlVHrnFgYm3VsNIzahSsvCD424AyvP', '': f'{gdrive_url}id=14X7FlkjrlUgmnsYtPwdh-gGuFla4yb5u', '': f'{gdrive_url}id=1ySNspxbK-snNoAZM7oxiWGvOnTRdSyEK' } } splits = { 'penn94': f'{github_url}/splits/fb100-Penn94-splits.npy', } def __init__(self, root: str, name: str, transform: Optional[Callable] = None, pre_transform: Optional[Callable] = None): = name.lower() assert in self.datasets.keys() super().__init__(root, transform, pre_transform), self.slices = torch.load(self.processed_paths[0]) @property def raw_dir(self) -> str: return osp.join(self.root,, 'raw') @property def processed_dir(self) -> str: return osp.join(self.root,, 'processed') @property def raw_file_names(self) -> List[str]: names = list(self.datasets[].keys()) if in self.splits: names += [self.splits[].split('/')[-1]] return names @property def processed_file_names(self) -> str: return '' def download(self): for filename, path in self.datasets[].items(): download_url(path, self.raw_dir, filename=filename) if in self.splits: download_url(self.splits[], self.raw_dir) def _process_wiki(self): paths = {x.split('/')[-1]: x for x in self.raw_paths} x = torch.load(paths['']) edge_index = torch.load(paths['']).t().contiguous() y = torch.load(paths['']) return Data(x=x, edge_index=edge_index, y=y) def _process_facebook(self): from import loadmat mat = loadmat(self.raw_paths[0]) A = mat['A'].tocsr().tocoo() row = torch.from_numpy(A.row).to(torch.long) col = torch.from_numpy(A.col).to(torch.long) edge_index = torch.stack([row, col], dim=0) metadata = torch.from_numpy(mat['local_info'].astype('int64')) xs = [] y = metadata[:, 1] - 1 # gender label, -1 means unlabeled x =[metadata[:, :1], metadata[:, 2:]], dim=-1) for i in range(x.size(1)): _, out = x[:, i].unique(return_inverse=True) xs.append(F.one_hot(out).to(torch.float)) x =, dim=-1) data = Data(x=x, edge_index=edge_index, y=y) if in self.splits: splits = np.load(self.raw_paths[1], allow_pickle=True) sizes = (data.num_nodes, len(splits)) data.train_mask = torch.zeros(sizes, dtype=torch.bool) data.val_mask = torch.zeros(sizes, dtype=torch.bool) data.test_mask = torch.zeros(sizes, dtype=torch.bool) for i, split in enumerate(splits): data.train_mask[:, i][torch.tensor(split['train'])] = True data.val_mask[:, i][torch.tensor(split['valid'])] = True data.test_mask[:, i][torch.tensor(split['test'])] = True return data def _process_genius(self): from import loadmat mat = loadmat(self.raw_paths[0]) edge_index = torch.from_numpy(mat['edge_index']).to(torch.long) x = torch.from_numpy(mat['node_feat']).to(torch.float) y = torch.from_numpy(mat['label']).squeeze().to(torch.long) return Data(x=x, edge_index=edge_index, y=y) def process(self): if in self.facebook_datasets: data = self._process_facebook() elif == 'genius': data = self._process_genius() elif == 'wiki': data = self._process_wiki() else: raise NotImplementedError( f"chosen dataset '{}' is not implemented") if self.pre_transform is not None: data = self.pre_transform(data)[data]), self.processed_paths[0]) def __repr__(self) -> str: return f'{}({len(self)})'