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

This commit is contained in:
Muhammed Fatih BALIN
2024-08-29 19:52:56 -04:00
committed by GitHub
parent b4bd509226
commit 647d8ea815
3 changed files with 14 additions and 10 deletions

View File

@@ -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

View File

@@ -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",

View File

@@ -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",