Source code for torch_geometric.datasets.webkb

import os.path as osp

import torch
import numpy as np
from torch_sparse import coalesce
from torch_geometric.data import InMemoryDataset, download_url, Data


[docs]class WebKB(InMemoryDataset): r"""The WebKB datasets used in the `"Geom-GCN: Geometric Graph Convolutional Networks" <https://openreview.net/forum?id=S1e2agrFvS>`_ paper. Nodes represent web pages and edges represent hyperlinks between them. Node features are the bag-of-words representation of web pages. The task is to classify the nodes into one of the five categories, student, project, course, staff, and faculty. Args: root (string): Root directory where the dataset should be saved. name (string): The name of the dataset (:obj:`"Cornell"`, :obj:`"Texas"`, :obj:`"Wisconsin"`). transform (callable, optional): A function/transform that takes in an :obj:`torch_geometric.data.Data` 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:`torch_geometric.data.Data` object and returns a transformed version. The data object will be transformed before being saved to disk. (default: :obj:`None`) """ url = 'https://raw.githubusercontent.com/graphdml-uiuc-jlu/geom-gcn/master' def __init__(self, root, name, transform=None, pre_transform=None): self.name = name.lower() assert self.name in ['cornell', 'texas', 'wisconsin'] super(WebKB, self).__init__(root, transform, pre_transform) self.data, self.slices = torch.load(self.processed_paths[0]) @property def raw_dir(self): return osp.join(self.root, self.name, 'raw') @property def processed_dir(self): return osp.join(self.root, self.name, 'processed') @property def raw_file_names(self): return ['out1_node_feature_label.txt', 'out1_graph_edges.txt'] + [ '{}_split_0.6_0.2_{}.npz'.format(self.name, i) for i in range(10) ] @property def processed_file_names(self): return 'data.pt' def download(self): for f in self.raw_file_names[:2]: download_url(f'{self.url}/new_data/{self.name}/{f}', self.raw_dir) for f in self.raw_file_names[2:]: download_url(f'{self.url}/splits/{f}', self.raw_dir) def process(self): with open(self.raw_paths[0], 'r') as f: data = f.read().split('\n')[1:-1] x = [[float(v) for v in r.split('\t')[1].split(',')] for r in data] x = torch.tensor(x, dtype=torch.float) y = [int(r.split('\t')[2]) for r in data] y = torch.tensor(y, dtype=torch.long) with open(self.raw_paths[1], 'r') as f: data = f.read().split('\n')[1:-1] data = [[int(v) for v in r.split('\t')] for r in data] edge_index = torch.tensor(data, dtype=torch.long).t().contiguous() edge_index, _ = coalesce(edge_index, None, x.size(0), x.size(0)) train_masks, val_masks, test_masks = [], [], [] for f in self.raw_paths[2:]: tmp = np.load(f) train_masks += [torch.from_numpy(tmp['train_mask']).to(torch.bool)] val_masks += [torch.from_numpy(tmp['val_mask']).to(torch.bool)] test_masks += [torch.from_numpy(tmp['test_mask']).to(torch.bool)] train_mask = torch.stack(train_masks, dim=1) val_mask = torch.stack(val_masks, dim=1) test_mask = torch.stack(test_masks, dim=1) data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask, val_mask=val_mask, test_mask=test_mask) data = data if self.pre_transform is None else self.pre_transform(data) torch.save(self.collate([data]), self.processed_paths[0]) def __repr__(self): return '{}()'.format(self.name)