Files
2026-07-13 13:35:51 +08:00

272 lines
8.2 KiB
Python

import os
import pytest
@pytest.mark.parametrize(
"data",
[
"cora",
"citeseer",
"pubmed",
"csv",
"reddit",
"co-buy-computer",
"ogbn-arxiv",
"ogbn-products",
],
)
def test_nodepred_data(data):
os.system(f"dgl configure nodepred --data {data} --model gcn")
assert os.path.exists(f"nodepred_{data}_gcn.yaml")
custom_cfg = f"custom_{data}_gcn.yaml"
os.system(
f"dgl configure nodepred --data {data} --model gcn --cfg {custom_cfg}"
)
assert os.path.exists(custom_cfg)
custom_script = f"{data}_gcn.py"
os.system(f"dgl export --cfg {custom_cfg} --output {custom_script}")
assert os.path.exists(custom_script)
@pytest.mark.parametrize("model", ["gcn", "gat", "sage", "sgc", "gin"])
def test_nodepred_model(model):
os.system(f"dgl configure nodepred --data cora --model {model}")
assert os.path.exists(f"nodepred_cora_{model}.yaml")
custom_cfg = f"custom_cora_{model}.yaml"
os.system(
f"dgl configure nodepred --data cora --model {model} --cfg {custom_cfg}"
)
assert os.path.exists(custom_cfg)
custom_script = f"cora_{model}.py"
os.system(f"dgl export --cfg {custom_cfg} --output {custom_script}")
assert os.path.exists(custom_script)
@pytest.mark.parametrize(
"data",
[
"cora",
"citeseer",
"pubmed",
"csv",
"reddit",
"co-buy-computer",
"ogbn-arxiv",
"ogbn-products",
],
)
def test_nodepred_ns_data(data):
os.system(f"dgl configure nodepred-ns --data {data} --model gcn")
assert os.path.exists(f"nodepred-ns_{data}_gcn.yaml")
custom_cfg = f"ns-custom_{data}_gcn.yaml"
os.system(
f"dgl configure nodepred-ns --data {data} --model gcn --cfg {custom_cfg}"
)
assert os.path.exists(custom_cfg)
custom_script = f"ns-{data}_gcn.py"
os.system(f"dgl export --cfg {custom_cfg} --output {custom_script}")
assert os.path.exists(custom_script)
@pytest.mark.parametrize("model", ["gcn", "gat", "sage"])
def test_nodepred_ns_model(model):
os.system(f"dgl configure nodepred-ns --data cora --model {model}")
assert os.path.exists(f"nodepred-ns_cora_{model}.yaml")
custom_cfg = f"ns-custom_cora_{model}.yaml"
os.system(
f"dgl configure nodepred-ns --data cora --model {model} --cfg {custom_cfg}"
)
assert os.path.exists(custom_cfg)
custom_script = f"ns-cora_{model}.py"
os.system(f"dgl export --cfg {custom_cfg} --output {custom_script}")
assert os.path.exists(custom_script)
@pytest.mark.parametrize(
"data",
[
"cora",
"citeseer",
"pubmed",
"csv",
"reddit",
"co-buy-computer",
"ogbn-arxiv",
"ogbn-products",
"ogbl-collab",
"ogbl-citation2",
],
)
def test_linkpred_data(data):
node_model = "gcn"
edge_model = "ele"
neg_sampler = "global"
custom_cfg = "_".join([data, node_model, edge_model, neg_sampler]) + ".yaml"
os.system(
"dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}".format(
data, node_model, edge_model, neg_sampler, custom_cfg
)
)
assert os.path.exists(custom_cfg)
custom_script = (
"_".join([data, node_model, edge_model, neg_sampler]) + ".py"
)
os.system(
"dgl export --cfg {} --output {}".format(custom_cfg, custom_script)
)
assert os.path.exists(custom_script)
@pytest.mark.parametrize("node_model", ["gcn", "gat", "sage", "sgc", "gin"])
def test_linkpred_node_model(node_model):
data = "cora"
edge_model = "ele"
neg_sampler = "global"
custom_cfg = "_".join([data, node_model, edge_model, neg_sampler]) + ".yaml"
os.system(
"dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}".format(
data, node_model, edge_model, neg_sampler, custom_cfg
)
)
assert os.path.exists(custom_cfg)
custom_script = (
"_".join([data, node_model, edge_model, neg_sampler]) + ".py"
)
os.system(
"dgl export --cfg {} --output {}".format(custom_cfg, custom_script)
)
assert os.path.exists(custom_script)
@pytest.mark.parametrize("edge_model", ["ele", "bilinear"])
def test_linkpred_edge_model(edge_model):
data = "cora"
node_model = "gcn"
neg_sampler = "global"
custom_cfg = "_".join([data, node_model, edge_model, neg_sampler]) + ".yaml"
os.system(
"dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}".format(
data, node_model, edge_model, neg_sampler, custom_cfg
)
)
assert os.path.exists(custom_cfg)
custom_script = (
"_".join([data, node_model, edge_model, neg_sampler]) + ".py"
)
os.system(
"dgl export --cfg {} --output {}".format(custom_cfg, custom_script)
)
assert os.path.exists(custom_script)
@pytest.mark.parametrize("neg_sampler", ["global", "persource", ""])
def test_linkpred_neg_sampler(neg_sampler):
data = "cora"
node_model = "gcn"
edge_model = "ele"
custom_cfg = f"{data}_{node_model}_{edge_model}_{neg_sampler}.yaml"
if neg_sampler == "":
os.system(
"dgl configure linkpred --data {} --node-model {} --edge-model {} --cfg {}".format(
data, node_model, edge_model, custom_cfg
)
)
else:
os.system(
"dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}".format(
data, node_model, edge_model, neg_sampler, custom_cfg
)
)
assert os.path.exists(custom_cfg)
custom_script = f"{data}_{node_model}_{edge_model}_{neg_sampler}.py"
os.system(
"dgl export --cfg {} --output {}".format(custom_cfg, custom_script)
)
assert os.path.exists(custom_script)
@pytest.mark.parametrize("data", ["csv", "ogbg-molhiv", "ogbg-molpcba"])
@pytest.mark.parametrize("model", ["gin", "pna"])
def test_graphpred(data, model):
os.system(
"dgl configure graphpred --data {} --model {}".format(data, model)
)
assert os.path.exists("graphpred_{}_{}.yaml".format(data, model))
custom_cfg = "custom_{}_{}.yaml".format(data, model)
os.system(
"dgl configure graphpred --data {} --model {} --cfg {}".format(
data, model, custom_cfg
)
)
assert os.path.exists(custom_cfg)
custom_script = "_".join([data, model]) + ".py"
os.system(
"dgl export --cfg {} --output {}".format(custom_cfg, custom_script)
)
assert os.path.exists(custom_script)
@pytest.mark.parametrize(
"recipe",
[
"graphpred_hiv_gin.yaml",
"graphpred_hiv_pna.yaml",
"graphpred_pcba_gin.yaml",
"linkpred_cora_sage.yaml",
"linkpred_citation2_sage.yaml",
"linkpred_collab_sage.yaml",
"nodepred_citeseer_gat.yaml",
"nodepred_citeseer_gcn.yaml",
"nodepred_citeseer_sage.yaml",
"nodepred_cora_gat.yaml",
"nodepred_cora_gcn.yaml",
"nodepred_cora_sage.yaml",
"nodepred_pubmed_gat.yaml",
"nodepred_pubmed_gcn.yaml",
"nodepred_pubmed_sage.yaml",
"nodepred-ns_arxiv_gcn.yaml",
"nodepred-ns_product_sage.yaml",
],
)
def test_recipe(recipe):
# Remove all generated yaml files
current_dir = os.listdir("./")
for item in current_dir:
if item.endswith(".yaml"):
os.remove(item)
os.system("dgl recipe get {}".format(recipe))
assert os.path.exists(recipe)
def test_node_cora():
os.system("dgl configure nodepred --data cora --model gcn")
os.system("dgl train --cfg nodepred_cora_gcn.yaml")
assert os.path.exists("results")
assert os.path.exists("results/run_0.pth")
os.system("dgl configure-apply nodepred --cpt results/run_0.pth")
assert os.path.exists("apply_nodepred_cora_gcn.yaml")
os.system(
"dgl configure-apply nodepred --data cora --cpt results/run_0.pth --cfg apply.yaml"
)
assert os.path.exists("apply.yaml")
os.system("dgl apply --cfg apply.yaml")
assert os.path.exists("apply_results/output.csv")
os.system("dgl export --cfg apply.yaml --output apply.py")
assert os.path.exists("apply.py")