Source code for torch_geometric.datasets.webkb

import os.path as osp

import torch
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:`"Washington"`, :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/new_data') def __init__(self, root, name, transform=None, pre_transform=None): self.name = name.lower() assert self.name in ['cornell', 'texas', 'washington', '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'] @property def processed_file_names(self): return 'data.pt' def download(self): for name in self.raw_file_names: download_url(f'{self.url}/{self.name}/{name}', 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.float) 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)) data = Data(x=x, edge_index=edge_index, y=y) 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)