mirror of
https://github.com/dmlc/dgl.git
synced 2026-06-03 19:34:33 +08:00
[GraphBolt][PyG] Add igb datasets to the examples. (#7758)
This commit is contained in:
committed by
GitHub
parent
b4bd509226
commit
647d8ea815
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user