[GraphBolt][PyG] Add igb datasets to the examples. (#7758)

dmlc · Aug 29, 2024 · 647d8ea · 647d8ea
1 parent b4bd509
commit 647d8ea
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/examples/graphbolt/pyg/labor/load_dataset.py b/examples/graphbolt/pyg/labor/load_dataset.py
@@ -39,12 +39,7 @@ def load_dataset(dataset_name, disk_based_feature_keys=None):
         "flickr",
     ]:
         dataset, multilabel = load_dgl(dataset_name)
-    elif dataset_name in [
-        "ogbn-products",
-        "ogbn-arxiv",
-        "ogbn-papers100M",
-        "ogbn-mag240M",
-    ]:
+    else:
         if "mag240M" in dataset_name:
             dataset_name = "ogb-lsc-mag240m"
         dataset = gb.BuiltinDataset(dataset_name)
@@ -56,7 +51,5 @@ def load_dataset(dataset_name, disk_based_feature_keys=None):
             if feature_key in disk_based_feature_keys:
                 feature["in_memory"] = False
         dataset = dataset.load()
-    else:
-        raise ValueError("unknown dataset")
 
     return dataset, multilabel
diff --git a/examples/graphbolt/pyg/labor/node_classification.py b/examples/graphbolt/pyg/labor/node_classification.py
@@ -363,6 +363,9 @@ def parse_args():
             "ogbn-arxiv",
             "ogbn-products",
             "ogbn-papers100M",
+            "igb-hom-tiny",
+            "igb-hom-small",
+            "igb-hom-medium",
             "reddit",
             "yelp",
             "flickr",

diff --git a/examples/graphbolt/pyg/node_classification_advanced.py b/examples/graphbolt/pyg/node_classification_advanced.py
@@ -335,9 +335,17 @@ def parse_args():
         "--dataset",
         type=str,
         default="ogbn-products",
-        choices=["ogbn-arxiv", "ogbn-products", "ogbn-papers100M"],
+        choices=[
+            "ogbn-arxiv",
+            "ogbn-products",
+            "ogbn-papers100M",
+            "igb-hom-tiny",
+            "igb-hom-small",
+            "igb-hom-medium",
+        ],
         help="The dataset we can use for node classification example. Currently"
-        " ogbn-products, ogbn-arxiv, ogbn-papers100M datasets are supported.",
+        " ogbn-products, ogbn-arxiv, ogbn-papers100M and"
+        " igb-hom-[tiny|small|medium] datasets are supported.",
     )
     parser.add_argument(
         "--fanout",