dataset_pyg.py

from torch_geometric.data import InMemoryDataset
import pandas as pd
import shutil, os
import os.path as osp
import torch
import numpy as np
from ogb.utils.url import decide_download, download_url, extract_zip
from ogb.io.read_graph_pyg import read_graph_pyg
import ogb
import re
from torch_geometric.data.separate import separate
import copy


class PygGraphPropPredDataset(InMemoryDataset):
    def __init__(self, name, root='dataset', processed_name='processed', transform=None, pre_transform=None,
                 meta_dict=None, mode='full'):
        '''
            - name (str): name of the dataset
            - root (str): root directory to store the dataset folder
            - transform, pre_transform (optional): transform/pre-transform graph objects
            - processed_name: type of processing for different nesting models
            - mode: 'full' for full loading, 'sequential' for loading one data one time
            - meta_dict: dictionary that stores all the meta-information about data. Default is None,
                    but when something is passed, it uses its information. Useful for debugging for external contributers.
        '''

        self.name = name  ## original name, e.g., ogbg-molhiv
        self.processed_name = processed_name if mode == 'full' else processed_name + '_seq'
        self.mode = mode
        if meta_dict is None:
            self.dir_name = '_'.join(name.split('-'))

            # check if previously-downloaded folder exists.
            # If so, use that one.
            if osp.exists(osp.join(root, self.dir_name + '_pyg')):
                self.dir_name = self.dir_name + '_pyg'

            self.original_root = root
            self.root = osp.join(root, self.dir_name)

            master = pd.read_csv(os.path.join(os.path.dirname(ogb.__file__), 'graphproppred', 'master.csv'),
                                 index_col=0)
            # master = pd.read_csv(os.path.join(os.path.dirname(__file__), 'master.csv'), index_col=0)
            if not self.name in master:
                error_mssg = 'Invalid dataset name {}.\n'.format(self.name)
                error_mssg += 'Available datasets are as follows:\n'
                error_mssg += '\n'.join(master.keys())
                raise ValueError(error_mssg)
            self.meta_info = master[self.name]

        else:
            self.dir_name = meta_dict['dir_path']
            self.original_root = ''
            self.root = meta_dict['dir_path']
            self.meta_info = meta_dict

        # check version
        # First check whether the dataset has been already downloaded or not.
        # If so, check whether the dataset version is the newest or not.
        # If the dataset is not the newest version, notify this to the user.
        if osp.isdir(self.root) and (
        not osp.exists(osp.join(self.root, 'RELEASE_v' + str(self.meta_info['version']) + '.txt'))):
            print(self.name + ' has been updated.')
            if input('Will you update the dataset now? (y/N)\n').lower() == 'y':
                shutil.rmtree(self.root)

        self.download_name = self.meta_info['download_name']  ## name of downloaded file, e.g., tox21

        self.num_tasks = int(self.meta_info['num tasks'])
        self.eval_metric = self.meta_info['eval metric']
        self.task_type = self.meta_info['task type']
        self.__num_classes__ = int(self.meta_info['num classes'])
        self.binary = self.meta_info['binary'] == 'True'

        super(PygGraphPropPredDataset, self).__init__(self.root, transform, pre_transform)

        if self.mode == 'full':
            self.data, self.slices = torch.load(self.processed_paths[0])
        elif self.mode == 'sequential':
            length = 0
            for file in os.listdir(self.processed_dir):
                try:
                    l = int(re.findall("\d+\.?\d*", file)[0][:-1])
                    if l > length:
                        length = l
                except:
                    continue
            self.data = [{'idx': i} for i in range(length + 1)]
            self.slices = {'idx': torch.tensor([i for i in range(length + 2)])}
        else:
            print('No such data loading mode!')
            exit(1)

    def get_idx_split(self, split_type=None):
        if split_type is None:
            split_type = self.meta_info['split']

        path = osp.join(self.root, 'split', split_type)

        # short-cut if split_dict.pt exists
        if os.path.isfile(os.path.join(path, 'split_dict.pt')):
            return torch.load(os.path.join(path, 'split_dict.pt'))

        train_idx = pd.read_csv(osp.join(path, 'train.csv.gz'), compression='gzip', header=None).values.T[0]
        valid_idx = pd.read_csv(osp.join(path, 'valid.csv.gz'), compression='gzip', header=None).values.T[0]
        test_idx = pd.read_csv(osp.join(path, 'test.csv.gz'), compression='gzip', header=None).values.T[0]

        return {'train': torch.tensor(train_idx, dtype=torch.long), 'valid': torch.tensor(valid_idx, dtype=torch.long),
                'test': torch.tensor(test_idx, dtype=torch.long)}

    @property
    def num_classes(self):
        return self.__num_classes__

    @property
    def raw_file_names(self):
        if self.binary:
            return ['data.npz']
        else:
            file_names = ['edge']
            if self.meta_info['has_node_attr'] == 'True':
                file_names.append('node-feat')
            if self.meta_info['has_edge_attr'] == 'True':
                file_names.append('edge-feat')
            return [file_name + '.csv.gz' for file_name in file_names]

    @property
    def processed_file_names(self):
        if self.mode == 'full':
            return 'geometric_data_processed.pt'
        else:
            return 'd0.pt'

    @property
    def processed_dir(self):
        return osp.join(self.root, self.processed_name)

    def download(self):
        url = self.meta_info['url']
        if decide_download(url):
            path = download_url(url, self.original_root)
            extract_zip(path, self.original_root)
            os.unlink(path)
            shutil.rmtree(self.root)
            shutil.move(osp.join(self.original_root, self.download_name), self.root)

        else:
            print('Stop downloading.')
            shutil.rmtree(self.root)
            exit(-1)

    def process(self):
        ### read pyg graph list
        add_inverse_edge = self.meta_info['add_inverse_edge'] == 'True'

        if self.meta_info['additional node files'] == 'None':
            additional_node_files = []
        else:
            additional_node_files = self.meta_info['additional node files'].split(',')

        if self.meta_info['additional edge files'] == 'None':
            additional_edge_files = []
        else:
            additional_edge_files = self.meta_info['additional edge files'].split(',')

        data_list = read_graph_pyg(self.raw_dir, add_inverse_edge=add_inverse_edge,
                                   additional_node_files=additional_node_files,
                                   additional_edge_files=additional_edge_files, binary=self.binary)

        if self.task_type == 'subtoken prediction':
            graph_label_notparsed = pd.read_csv(osp.join(self.raw_dir, 'graph-label.csv.gz'), compression='gzip',
                                                header=None).values
            graph_label = [str(graph_label_notparsed[i][0]).split(' ') for i in range(len(graph_label_notparsed))]

            for i, g in enumerate(data_list):
                g.y = graph_label[i]

        else:
            if self.binary:
                graph_label = np.load(osp.join(self.raw_dir, 'graph-label.npz'))['graph_label']
            else:
                graph_label = pd.read_csv(osp.join(self.raw_dir, 'graph-label.csv.gz'), compression='gzip',
                                          header=None).values

            has_nan = np.isnan(graph_label).any()

            for i, g in enumerate(data_list):
                if 'classification' in self.task_type:
                    if has_nan:
                        g.y = torch.from_numpy(graph_label[i]).view(1, -1).to(torch.float32)
                    else:
                        g.y = torch.from_numpy(graph_label[i]).view(1, -1).to(torch.long)
                else:
                    g.y = torch.from_numpy(graph_label[i]).view(1, -1).to(torch.float32)

        if self.mode == 'full':
            if self.pre_transform is not None:
                # data_list = [self.pre_transform(data) for data in data_list]
                data_list_new = []
                for i, data in enumerate(data_list):
                    if i % 1000 == 0:
                        print('Pre-processing: %d/%d' %(i, len(data_list)))
                    data_list_new.append(self.pre_transform(data))
                data_list = data_list_new
            data, slices = self.collate(data_list)
            print('Saving full dataset...')
            torch.save((data, slices), self.processed_paths[0])
        elif self.mode == 'sequential':
            print('Processing and saving sequential dataset...')
            if not osp.exists(self.processed_dir):
                os.makedirs(self.processed_dir)
            for i, d in enumerate(data_list):
                if i % 1000 == 0:
                    print('Pre-processing: %d/%d' % (i, len(data_list)))
                if self.pre_transform is not None:
                    d = self.pre_transform(d)
                data_save_path = osp.join(self.processed_dir, 'd%d.pt'%i)
                torch.save(d, data_save_path)
        else:
            print('Saving mode error: no such saving format!')
            exit(1)

    def get(self, idx: int):
        if self.mode == 'full':
            if self.len() == 1:
                return copy.copy(self.data)

            if not hasattr(self, '_data_list') or self._data_list is None:
                self._data_list = self.len() * [None]
            elif self._data_list[idx] is not None:
                return copy.copy(self._data_list[idx])

            data = separate(
                cls=self.data.__class__,
                batch=self.data,
                idx=idx,
                slice_dict=self.slices,
                decrement=False,
            )

            self._data_list[idx] = copy.copy(data)
        else:
            data = torch.load(osp.join(self.processed_dir, 'd%d.pt'%self.data[idx]['idx']))
        return data


if __name__ == '__main__':
    # pyg_dataset = PygGraphPropPredDataset(name = 'ogbg-molpcba')
    # print(pyg_dataset.num_classes)
    # split_index = pyg_dataset.get_idx_split()
    # print(pyg_dataset)
    # print(pyg_dataset[0])
    # print(pyg_dataset[0].y)
    # print(pyg_dataset[0].y.dtype)
    # print(pyg_dataset[0].edge_index)
    # print(pyg_dataset[split_index['train']])
    # print(pyg_dataset[split_index['valid']])
    # print(pyg_dataset[split_index['test']])

    pyg_dataset = PygGraphPropPredDataset(name='ogbg-code2')
    print(pyg_dataset.num_classes)
    split_index = pyg_dataset.get_idx_split()
    print(pyg_dataset[0])
    # print(pyg_dataset[0].node_is_attributed)
    print([pyg_dataset[i].x[1] for i in range(100)])
    # print(pyg_dataset[0].y)
    # print(pyg_dataset[0].edge_index)
    print(pyg_dataset[split_index['train']])
    print(pyg_dataset[split_index['valid']])
    print(pyg_dataset[split_index['test']])

    # from torch_geometric.loader import DataLoader
    # loader = DataLoader(pyg_dataset, batch_size=32, shuffle=True)
    # for batch in loader:
    #     print(batch)
    #     print(batch.y)
    #     print(len(batch.y))

    #     exit(-1)