Merge pull request #60 from snap-stanford/external_dataset

DatasetSaver class, binary loading and saving, and DGL v0.5

Author: rusty1s

.gitignore

@@ -13,3 +13,8 @@ dataset/
 *.sh
 *analyze*
 *random.py
+*RELEASE_*
+*.csv.gz
+*.zip
+*submission_
+*.npz

README.md

@@ -1,5 +1,5 @@
-<p align="center">
-  <img width="40%" src="https://snap-stanford.github.io/ogb-web/assets/img/OGB_rectangle.png" />
+<p align='center'>
+  <img width='40%' src='https://snap-stanford.github.io/ogb-web/assets/img/OGB_rectangle.png' />
 </p>
 
 --------------------------------------------------------------------------------
@@ -12,8 +12,8 @@
 The Open Graph Benchmark (OGB) is a collection of benchmark datasets, data loaders, and evaluators for graph machine learning. Datasets cover a variety of graph machine learning tasks and real-world applications.
 The OGB data loaders are fully compatible with popular graph deep learning frameworks, including [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/) and [Deep Graph Library (DGL)](https://www.dgl.ai/). They provide automatic dataset downloading, standardized dataset splits, and unified performance evaluation.
 
-<p align="center">
-  <img width="80%" src="https://snap-stanford.github.io/ogb-web/assets/img/ogb_overview.png" />
+<p align='center'>
+  <img width='80%' src='https://snap-stanford.github.io/ogb-web/assets/img/ogb_overview.png' />
 </p>
 
 OGB aims to provide graph datasets that cover important graph machine learning tasks, diverse dataset scale, and rich domains.
@@ -24,8 +24,8 @@ OGB aims to provide graph datasets that cover important graph machine learning t
 
 **Rich domains:** Graph datasets come from diverse domains ranging from scientific ones to social/information networks, and also include heterogeneous knowledge graphs. 
 
-<p align="center">
-  <img width="70%" src="https://snap-stanford.github.io/ogb-web/assets/img/dataset_overview.png" />
+<p align='center'>
+  <img width='70%' src='https://snap-stanford.github.io/ogb-web/assets/img/dataset_overview.png' />
 </p>
 
 OGB is an on-going effort, and we are planning to increase our coverage in the future.
@@ -38,7 +38,7 @@ The release note is available [here](https://github.com/snap-stanford/ogb/releas
 #### Requirements
  - Python>=3.5
  - PyTorch>=1.2
- - DGL>=0.4.1 or torch-geometric>=1.3.1
+ - DGL>=0.5.0 or torch-geometric>=1.6.0
  - Numpy>=1.16.0
  - pandas>=0.24.0
  - urllib3>=1.24.0
@@ -77,12 +77,12 @@ Below, on PyTorch Geometric, we see that a few lines of code is sufficient to pr
 from ogb.graphproppred import PygGraphPropPredDataset
 from torch_geometric.data import DataLoader
 
-dataset = PygGraphPropPredDataset(name = "ogbg-molhiv")
+dataset = PygGraphPropPredDataset(name = 'ogbg-molhiv')
 
 split_idx = dataset.get_idx_split() 
-train_loader = DataLoader(dataset[split_idx["train"]], batch_size=32, shuffle=True)
-valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=32, shuffle=False)
-test_loader = DataLoader(dataset[split_idx["test"]], batch_size=32, shuffle=False)
+train_loader = DataLoader(dataset[split_idx['train']], batch_size=32, shuffle=True)
+valid_loader = DataLoader(dataset[split_idx['valid']], batch_size=32, shuffle=False)
+test_loader = DataLoader(dataset[split_idx['test']], batch_size=32, shuffle=False)
 ```
 
 #### (2) Evaluators
@@ -91,12 +91,12 @@ The standardized evaluation protocol allows researchers to reliably compare thei
 ```python
 from ogb.graphproppred import Evaluator
 
-evaluator = Evaluator(name = "ogbg-molhiv")
+evaluator = Evaluator(name = 'ogbg-molhiv')
 # You can learn the input and output format specification of the evaluator as follows.
 # print(evaluator.expected_input_format) 
 # print(evaluator.expected_output_format) 
-input_dict = {"y_true": y_true, "y_pred": y_pred}
-result_dict = evaluator.eval(input_dict) # E.g., {"rocauc": 0.7321}
+input_dict = {'y_true': y_true, 'y_pred': y_pred}
+result_dict = evaluator.eval(input_dict) # E.g., {'rocauc': 0.7321}
 ```
 
 ## Citing OGB

ogb/graphproppred/dataset.py

@@ -3,40 +3,56 @@
 import numpy as np
 import os.path as osp
 from ogb.utils.url import decide_download, download_url, extract_zip
-from ogb.io.read_graph_raw import read_csv_graph_raw
+from ogb.io.read_graph_raw import read_csv_graph_raw, read_binary_graph_raw
 import torch
 
 class GraphPropPredDataset(object):
-    def __init__(self, name, root = "dataset"):
-        self.name = name ## original name, e.g., ogbg-mol-tox21
-        self.dir_name = "_".join(name.split("-")) ## replace hyphen with underline, e.g., ogbg_mol_tox21
-
-        self.original_root = root
-        self.root = osp.join(root, self.dir_name)
-
-        self.meta_info = pd.read_csv(os.path.join(os.path.dirname(__file__), "master.csv"), index_col = 0)
-        if not self.name in self.meta_info:
-            print(self.name)
-            error_mssg = "Invalid dataset name {}.\n".format(self.name)
-            error_mssg += "Available datasets are as follows:\n"
-            error_mssg += "\n".join(self.meta_info.keys())
-            raise ValueError(error_mssg)
+    def __init__(self, name, root = 'dataset', meta_dict = None):
+        '''
+            - name (str): name of the dataset
+            - root (str): root directory to store the dataset folder
+
+            - 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-hib
+        
+        if meta_dict is None:
+            self.dir_name = '_'.join(name.split('-')) ## replace hyphen with underline, e.g., ogbg_hiv
+            self.original_root = root
+            self.root = osp.join(root, self.dir_name)
+            
+            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[self.name]['version']) + '.txt'))):
+        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":
+            if input('Will you update the dataset now? (y/N)\n').lower() == 'y':
                 shutil.rmtree(self.root)
 
-        self.download_name = self.meta_info[self.name]["download_name"] ## name of downloaded file, e.g., tox21
+        self.download_name = self.meta_info['download_name'] ## name of downloaded file, e.g., tox21
 
-        self.num_tasks = int(self.meta_info[self.name]["num tasks"])
-        self.eval_metric = self.meta_info[self.name]["eval metric"]
-        self.task_type = self.meta_info[self.name]["task type"]
-        self.num_classes = self.meta_info[self.name]["num classes"]
+        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 = self.meta_info['num classes']
+        self.binary = self.meta_info['binary'] == 'True'
 
         super(GraphPropPredDataset, self).__init__()
 
@@ -52,63 +68,81 @@ def pre_process(self):
             self.graphs, self.labels = loaded_dict['graphs'], loaded_dict['labels']
 
         else:
-            ### download
-            url = self.meta_info[self.name]["url"]
-            if decide_download(url):
-                path = download_url(url, self.original_root)
-                extract_zip(path, self.original_root)
-                os.unlink(path)
-                # delete folder if there exists
-                try:
-                    shutil.rmtree(self.root)
-                except:
-                    pass
-                shutil.move(osp.join(self.original_root, self.download_name), self.root)
+            ### check download
+            if self.binary:
+                # npz format
+                has_necessary_file = osp.exists(osp.join(self.root, 'raw', 'data.npz'))
             else:
-                print("Stop download.")
-                exit(-1)
+                # csv file
+                has_necessary_file = osp.exists(osp.join(self.root, 'raw', 'edge.csv.gz'))
+            
+            ### download
+            if not has_necessary_file:
+                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)
+                    # delete folder if there exists
+                    try:
+                        shutil.rmtree(self.root)
+                    except:
+                        pass
+                    shutil.move(osp.join(self.original_root, self.download_name), self.root)
+                else:
+                    print('Stop download.')
+                    exit(-1)
 
             ### preprocess
-            add_inverse_edge = self.meta_info[self.name]["add_inverse_edge"] == "True"
+            add_inverse_edge = self.meta_info['add_inverse_edge'] == 'True'
 
-            if self.meta_info[self.name]["additional node files"] == 'None':
+            if self.meta_info['additional node files'] == 'None':
                 additional_node_files = []
             else:
-                additional_node_files = self.meta_info[self.name]["additional node files"].split(',')
+                additional_node_files = self.meta_info['additional node files'].split(',')
 
-            if self.meta_info[self.name]["additional edge files"] == 'None':
+            if self.meta_info['additional edge files'] == 'None':
                 additional_edge_files = []
             else:
-                additional_edge_files = self.meta_info[self.name]["additional edge files"].split(',')
-
-            self.graphs = read_csv_graph_raw(raw_dir, add_inverse_edge = add_inverse_edge, additional_node_files = additional_node_files, additional_edge_files = additional_edge_files)
-
+                additional_edge_files = self.meta_info['additional edge files'].split(',')
+            
+            if self.binary:
+                self.graphs = read_binary_graph_raw(raw_dir, add_inverse_edge = add_inverse_edge)
+            else:
+                self.graphs = read_csv_graph_raw(raw_dir, add_inverse_edge = add_inverse_edge, additional_node_files = additional_node_files, additional_edge_files = additional_edge_files)
 
             if self.task_type == 'subtoken prediction':
-                labels_joined = pd.read_csv(osp.join(raw_dir, "graph-label.csv.gz"), compression="gzip", header = None).values
+                labels_joined = pd.read_csv(osp.join(raw_dir, 'graph-label.csv.gz'), compression='gzip', header = None).values
                 # need to split each element into subtokens
                 self.labels = [str(labels_joined[i][0]).split(' ') for i in range(len(labels_joined))]
             else:
-                self.labels = pd.read_csv(osp.join(raw_dir, "graph-label.csv.gz"), compression="gzip", header = None).values
+                if self.binary:
+                    self.labels = np.load(osp.join(raw_dir, 'graph-label.npz'))['graph_label']
+                else:
+                    self.labels = pd.read_csv(osp.join(raw_dir, 'graph-label.csv.gz'), compression='gzip', header = None).values
 
             print('Saving...')
             torch.save({'graphs': self.graphs, 'labels': self.labels}, pre_processed_file_path, pickle_protocol=4)
 
 
     def get_idx_split(self, split_type = None):
         if split_type is None:
-            split_type = self.meta_info[self.name]["split"]
+            split_type = self.meta_info['split']
 
-        path = osp.join(self.root, "split", split_type)
+        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]
+        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": train_idx, "valid": valid_idx, "test": test_idx}
+        return {'train': train_idx, 'valid': valid_idx, 'test': test_idx}
 
     def __getitem__(self, idx):
-        """Get datapoint with index"""
+        '''Get datapoint with index'''
 
         if isinstance(idx, (int, np.integer)):
             return self.graphs[idx], self.labels[idx]
@@ -117,20 +151,20 @@ def __getitem__(self, idx):
             'Only integer is valid index (got {}).'.format(type(idx).__name__))
 
     def __len__(self):
-        """Length of the dataset
+        '''Length of the dataset
         Returns
         -------
         int
             Length of Dataset
-        """
+        '''
         return len(self.graphs)
 
     def __repr__(self):  # pragma: no cover
         return '{}({})'.format(self.__class__.__name__, len(self))
 
 
-if __name__ == "__main__":
-    dataset = GraphPropPredDataset(name = "ogbg-code")
+if __name__ == '__main__':
+    dataset = GraphPropPredDataset(name = 'ogbg-code')
     # target_list = np.array([len(label) for label in dataset.labels])
     # print(np.sum(target_list == 1)/ float(len(target_list)))
     # print(np.sum(target_list == 2)/ float(len(target_list)))
@@ -144,8 +178,8 @@ def __repr__(self):  # pragma: no cover
     print(split_index)
     # print(dataset)
     # print(dataset[2])
-    # print(split_index["train"])
-    # print(split_index["valid"])
-    # print(split_index["test"])
+    # print(split_index['train'])
+    # print(split_index['valid'])
+    # print(split_index['test'])