272 lines
8.2 KiB
Python
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")
|