Fix OGB load data

gdmnl · May 21, 2024 · d9b48b7 · d9b48b7
1 parent c051443
commit d9b48b7
Show file tree

Hide file tree

Showing 10 changed files with 1,283 additions and 1,001 deletions.
diff --git a/examples/dataset_process/__init__.py b/examples/dataset_process/__init__.py
@@ -4,5 +4,6 @@
 from .utils import (
     idx2mask,
     split_random,
-    even_quantile_labels
+    even_quantile_labels,
+    get_iso_nodes_mapping
 )
diff --git a/examples/dataset_process/linkx.py b/examples/dataset_process/linkx.py
@@ -203,3 +203,13 @@ def forward(self, data: Any) -> Any:
         data.y = torch.tensor(y, dtype=torch.long)
         del data['node_year']
         return data
+
+
+class T_ogbn_mag(BaseTransform):
+    def forward(self, data: Any) -> Any:
+        new_data = Data(
+            x=data.x_dict['paper'],
+            edge_index=data.edge_index_dict[('paper', 'cites', 'paper')],
+            y=data.y_dict['paper'],
+            num_nodes=data.x_dict['paper'].shape[0])
+        return new_data
diff --git a/examples/dataset_process/utils.py b/examples/dataset_process/utils.py
@@ -57,3 +57,16 @@ def even_quantile_labels(vals, nclasses, verbose=True):
         for class_idx, interval in enumerate(interval_lst):
             print(f'Class {class_idx}: [{interval[0]}, {interval[1]})]')
     return label
+
+
+def get_iso_nodes_mapping(dataset):
+    data = dataset.get(dataset.indices()[0])
+    edge_index = data.edge_index
+    src, dst = edge_index[0], edge_index[1]
+    bin = torch.zeros(data.num_nodes, dtype=torch.bool)
+    bin[src] = True
+    bin[dst] = True
+    kept_nodes = torch.where(bin)[0]
+    mapping = torch.zeros(data.num_nodes, dtype=torch.long) - 1
+    mapping[kept_nodes] = torch.arange(kept_nodes.shape[0])
+    return mapping
diff --git a/examples/scripts/pfb_bank_s.sh b/examples/scripts/pfb_bank_s.sh
@@ -29,7 +29,9 @@ ARGS_S=(
     "--suffix" "fb_bank"
 )
 
-DATAS=("cora" "citeseer" "pubmed" "flickr" "chameleon_filtered" "squirrel_filtered" "actor" "roman_empire")
+# DATAS=("cora" "citeseer" "pubmed" "flickr" "chameleon_filtered" "squirrel_filtered" "actor" "roman_empire")
+# DATAS=("amazon_ratings" "minesweeper" "tolokers" "questions" "reddit" "penn94")
+DATAS=("ogbn-arxiv" "arxiv-year" "genius" "twitch-gamer" "ogbn-mag" "pokec")
 
 for data in ${DATAS[@]}; do
     PARLIST="normg,dp_lin,dp_conv,lr_lin,lr_conv,wd_lin,wd_conv"

diff --git a/examples/scripts/pfb_fix_s.sh b/examples/scripts/pfb_fix_s.sh
@@ -27,7 +27,10 @@ ARGS_S=(
     "--suffix" "fb_fix"
 )
 
-DATAS=("cora" "citeseer" "pubmed" "flickr" "chameleon_filtered" "squirrel_filtered" "actor" "roman_empire")
+# DATAS=("cora" "citeseer" "pubmed" "flickr" "chameleon_filtered" "squirrel_filtered" "actor" "roman_empire")
+# DATAS=("amazon_ratings" "minesweeper" "tolokers" "questions" "reddit" "penn94")
+# DATAS=("ogbn-arxiv" "arxiv-year" "genius" "twitch-gamer" "ogbn-mag" "pokec")
+DATAS=("twitch-gamer" "ogbn-mag" "pokec")
 MODELS=("DecoupledFixed")
 CONVS=AdjConv
 SCHEMES=("impulse" "mono" "appr" "hk" "gaussian")

diff --git a/examples/scripts/pfb_var_s.sh b/examples/scripts/pfb_var_s.sh
@@ -31,7 +31,9 @@ ARGS_S=(
     "--suffix" "fb_var"
 )
 
-DATAS=("cora" "citeseer" "pubmed" "flickr" "chameleon_filtered" "squirrel_filtered" "actor" "roman_empire")
+# DATAS=("cora" "citeseer" "pubmed" "flickr" "chameleon_filtered" "squirrel_filtered" "actor" "roman_empire")
+# DATAS=("amazon_ratings" "minesweeper" "tolokers" "questions" "reddit" "penn94")
+DATAS=("ogbn-arxiv" "arxiv-year" "genius" "twitch-gamer" "ogbn-mag" "pokec")
 MODELS=("DecoupledVar")
 CONVS=("AdjConv" "ChebConv" "ChebConv2" "BernConv" "ClenhawConv" "HornerConv")