Skip to content

GitLab

Unverified Commit dce89919 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Auto-reformat multiple python folders. (#5325) 



* auto-reformat

* lintrunner

---------
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
parent ab812179
Hide whitespace changes
Inline Side-by-side
Showing with 456 additions and 242 deletions
dglgo/dglgo/apply_pipeline/__init__.py
View file @ dce89919
from .graphpred import ApplyGraphpredPipeline
from .nodepred import ApplyNodepredPipeline
from .nodepred_sample import ApplyNodepredNsPipeline
from .graphpred import ApplyGraphpredPipeline
dglgo/dglgo/apply_pipeline/graphpred/gen.py
View file @ dce89919
from copy import deepcopy
from pathlib import Path
from typing import Optional
import ruamel.yaml
import torch
import typer
from copy import deepcopy
from jinja2 import Template
from pathlib import Path
from pydantic import BaseModel, Field
from typing import Optional
from ...utils.factory import ApplyPipelineFactory, PipelineBase, DataFactory, GraphModelFactory
from ...utils.factory import (
ApplyPipelineFactory,
DataFactory,
GraphModelFactory,
PipelineBase,
)
from ...utils.yaml_dump import deep_convert_dict, merge_comment
pipeline_comments = {
"batch_size": "Graph batch size",
"num_workers": "Number of workers for data loading",
"save_path": "Directory to save the inference results"
"save_path": "Directory to save the inference results",
}
class ApplyGraphpredPipelineCfg(BaseModel):
batch_size: int = 32
num_workers: int = 4
save_path: str = "apply_results"
@ApplyPipelineFactory.register("graphpred")
class ApplyGraphpredPipeline(PipelineBase):
def __init__(self):
self.pipeline = {
"name": "graphpred",
"mode": "apply"
}
self.pipeline = {"name": "graphpred", "mode": "apply"}
@classmethod
def setup_user_cfg_cls(cls):
from ...utils.enter_config import UserConfig
class ApplyGraphPredUserConfig(UserConfig):
data: DataFactory.filter("graphpred").get_pydantic_config() = Field(..., discriminator="name")
general_pipeline: ApplyGraphpredPipelineCfg = ApplyGraphpredPipelineCfg()
data: DataFactory.filter("graphpred").get_pydantic_config() = Field(
..., discriminator="name"
)
general_pipeline: ApplyGraphpredPipelineCfg = (
ApplyGraphpredPipelineCfg()
)
cls.user_cfg_cls = ApplyGraphPredUserConfig
......@@ -45,9 +54,13 @@ class ApplyGraphpredPipeline(PipelineBase):
def get_cfg_func(self):
def config(
data: DataFactory.filter("graphpred").get_dataset_enum() = typer.Option(None, help="input data name"),
cfg: Optional[str] = typer.Option(None, help="output configuration file path"),
cpt: str = typer.Option(..., help="input checkpoint file path")
data: DataFactory.filter(
"graphpred"
).get_dataset_enum() = typer.Option(None, help="input data name"),
cfg: Optional[str] = typer.Option(
None, help="output configuration file path"
),
cpt: str = typer.Option(..., help="input checkpoint file path"),
):
# Training configuration
train_cfg = torch.load(cpt)["cfg"]
......@@ -57,7 +70,12 @@ class ApplyGraphpredPipeline(PipelineBase):
else:
data = data.name
if cfg is None:
cfg = "_".join(["apply", "graphpred", data, train_cfg["model_name"]]) + ".yaml"
cfg = (
"_".join(
["apply", "graphpred", data, train_cfg["model_name"]]
)
+ ".yaml"
)
self.__class__.setup_user_cfg_cls()
generated_cfg = {
......@@ -66,23 +84,31 @@ class ApplyGraphpredPipeline(PipelineBase):
"device": train_cfg["device"],
"data": {"name": data},
"cpt_path": cpt,
"general_pipeline": {"batch_size": train_cfg["general_pipeline"]["eval_batch_size"],
"num_workers": train_cfg["general_pipeline"]["num_workers"]}
"general_pipeline": {
"batch_size": train_cfg["general_pipeline"][
"eval_batch_size"
],
"num_workers": train_cfg["general_pipeline"]["num_workers"],
},
}
output_cfg = self.user_cfg_cls(**generated_cfg).dict()
output_cfg = deep_convert_dict(output_cfg)
# Not applicable for inference
output_cfg['data'].pop('split_ratio')
output_cfg["data"].pop("split_ratio")
comment_dict = {
"device": "Torch device name, e.g., cpu or cuda or cuda:0",
"cpt_path": "Path to the checkpoint file",
"general_pipeline": pipeline_comments
"general_pipeline": pipeline_comments,
}
comment_dict = merge_comment(output_cfg, comment_dict)
yaml = ruamel.yaml.YAML()
yaml.dump(comment_dict, Path(cfg).open("w"))
print("Configuration file is generated at {}".format(Path(cfg).absolute()))
print(
"Configuration file is generated at {}".format(
Path(cfg).absolute()
)
)
return config
......@@ -100,8 +126,12 @@ class ApplyGraphpredPipeline(PipelineBase):
model_name = train_cfg["model_name"]
model_code = GraphModelFactory.get_source_code(model_name)
render_cfg["model_code"] = model_code
render_cfg["model_class_name"] = GraphModelFactory.get_model_class_name(model_name)
render_cfg.update(DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"]))
render_cfg["model_class_name"] = GraphModelFactory.get_model_class_name(
model_name
)
render_cfg.update(
DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"])
)
# Dict for defining cfg in the rendered code
generated_user_cfg = deepcopy(user_cfg_dict)
......
dglgo/dglgo/apply_pipeline/nodepred/gen.py
View file @ dce89919
from copy import deepcopy
from pathlib import Path
from typing import Optional
import ruamel.yaml
import torch
import typer
from copy import deepcopy
from jinja2 import Template
from pathlib import Path
from pydantic import Field
from typing import Optional
from ...utils.factory import ApplyPipelineFactory, PipelineBase, DataFactory, NodeModelFactory
from ...utils.factory import (
ApplyPipelineFactory,
DataFactory,
NodeModelFactory,
PipelineBase,
)
from ...utils.yaml_dump import deep_convert_dict, merge_comment
@ApplyPipelineFactory.register("nodepred")
class ApplyNodepredPipeline(PipelineBase):
def __init__(self):
self.pipeline = {
"name": "nodepred",
"mode": "apply"
}
self.pipeline = {"name": "nodepred", "mode": "apply"}
@classmethod
def setup_user_cfg_cls(cls):
from ...utils.enter_config import UserConfig
class ApplyNodePredUserConfig(UserConfig):
data: DataFactory.filter("nodepred").get_pydantic_config() = Field(..., discriminator="name")
data: DataFactory.filter("nodepred").get_pydantic_config() = Field(
..., discriminator="name"
)
cls.user_cfg_cls = ApplyNodePredUserConfig
......@@ -34,9 +39,13 @@ class ApplyNodepredPipeline(PipelineBase):
def get_cfg_func(self):
def config(
data: DataFactory.filter("nodepred").get_dataset_enum() = typer.Option(None, help="input data name"),
cfg: Optional[str] = typer.Option(None, help="output configuration file path"),
cpt: str = typer.Option(..., help="input checkpoint file path")
data: DataFactory.filter(
"nodepred"
).get_dataset_enum() = typer.Option(None, help="input data name"),
cfg: Optional[str] = typer.Option(
None, help="output configuration file path"
),
cpt: str = typer.Option(..., help="input checkpoint file path"),
):
# Training configuration
train_cfg = torch.load(cpt)["cfg"]
......@@ -46,7 +55,12 @@ class ApplyNodepredPipeline(PipelineBase):
else:
data = data.name
if cfg is None:
cfg = "_".join(["apply", "nodepred", data, train_cfg["model_name"]]) + ".yaml"
cfg = (
"_".join(
["apply", "nodepred", data, train_cfg["model_name"]]
)
+ ".yaml"
)
self.__class__.setup_user_cfg_cls()
generated_cfg = {
......@@ -55,22 +69,28 @@ class ApplyNodepredPipeline(PipelineBase):
"device": train_cfg["device"],
"data": {"name": data},
"cpt_path": cpt,
"general_pipeline": {"save_path": "apply_results"}
"general_pipeline": {"save_path": "apply_results"},
}
output_cfg = self.user_cfg_cls(**generated_cfg).dict()
output_cfg = deep_convert_dict(output_cfg)
# Not applicable for inference
output_cfg['data'].pop('split_ratio')
output_cfg["data"].pop("split_ratio")
comment_dict = {
"device": "Torch device name, e.g., cpu or cuda or cuda:0",
"cpt_path": "Path to the checkpoint file",
"general_pipeline": {"save_path": "Directory to save the inference results"}
"general_pipeline": {
"save_path": "Directory to save the inference results"
},
}
comment_dict = merge_comment(output_cfg, comment_dict)
yaml = ruamel.yaml.YAML()
yaml.dump(comment_dict, Path(cfg).open("w"))
print("Configuration file is generated at {}".format(Path(cfg).absolute()))
print(
"Configuration file is generated at {}".format(
Path(cfg).absolute()
)
)
return config
......@@ -88,8 +108,12 @@ class ApplyNodepredPipeline(PipelineBase):
model_name = train_cfg["model_name"]
model_code = NodeModelFactory.get_source_code(model_name)
render_cfg["model_code"] = model_code
render_cfg["model_class_name"] = NodeModelFactory.get_model_class_name(model_name)
render_cfg.update(DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"]))
render_cfg["model_class_name"] = NodeModelFactory.get_model_class_name(
model_name
)
render_cfg.update(
DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"])
)
# Dict for defining cfg in the rendered code
generated_user_cfg = deepcopy(user_cfg_dict)
......
dglgo/dglgo/apply_pipeline/nodepred_sample/gen.py
View file @ dce89919
from copy import deepcopy
from pathlib import Path
from typing import Optional
import ruamel.yaml
import typer
import torch
from copy import deepcopy
import typer
from jinja2 import Template
from pathlib import Path
from pydantic import Field
from typing import Optional
from ...utils.factory import ApplyPipelineFactory, PipelineBase, DataFactory, NodeModelFactory
from ...utils.factory import (
ApplyPipelineFactory,
DataFactory,
NodeModelFactory,
PipelineBase,
)
from ...utils.yaml_dump import deep_convert_dict, merge_comment
@ApplyPipelineFactory.register("nodepred-ns")
class ApplyNodepredNsPipeline(PipelineBase):
def __init__(self):
self.pipeline = {
"name": "nodepred-ns",
"mode": "apply"
}
self.pipeline = {"name": "nodepred-ns", "mode": "apply"}
@classmethod
def setup_user_cfg_cls(cls):
from ...utils.enter_config import UserConfig
class ApplyNodePredUserConfig(UserConfig):
data: DataFactory.filter("nodepred-ns").get_pydantic_config() = Field(..., discriminator="name")
data: DataFactory.filter(
"nodepred-ns"
).get_pydantic_config() = Field(..., discriminator="name")
cls.user_cfg_cls = ApplyNodePredUserConfig
......@@ -34,9 +39,13 @@ class ApplyNodepredNsPipeline(PipelineBase):
def get_cfg_func(self):
def config(
data: DataFactory.filter("nodepred-ns").get_dataset_enum() = typer.Option(None, help="input data name"),
cfg: Optional[str] = typer.Option(None, help="output configuration file path"),
cpt: str = typer.Option(..., help="input checkpoint file path")
data: DataFactory.filter(
"nodepred-ns"
).get_dataset_enum() = typer.Option(None, help="input data name"),
cfg: Optional[str] = typer.Option(
None, help="output configuration file path"
),
cpt: str = typer.Option(..., help="input checkpoint file path"),
):
# Training configuration
train_cfg = torch.load(cpt)["cfg"]
......@@ -46,7 +55,12 @@ class ApplyNodepredNsPipeline(PipelineBase):
else:
data = data.name
if cfg is None:
cfg = "_".join(["apply", "nodepred-ns", data, train_cfg["model_name"]]) + ".yaml"
cfg = (
"_".join(
["apply", "nodepred-ns", data, train_cfg["model_name"]]
)
+ ".yaml"
)
self.__class__.setup_user_cfg_cls()
generated_cfg = {
......@@ -55,22 +69,28 @@ class ApplyNodepredNsPipeline(PipelineBase):
"device": train_cfg["device"],
"data": {"name": data},
"cpt_path": cpt,
"general_pipeline": {"save_path": "apply_results"}
"general_pipeline": {"save_path": "apply_results"},
}
output_cfg = self.user_cfg_cls(**generated_cfg).dict()
output_cfg = deep_convert_dict(output_cfg)
# Not applicable for inference
output_cfg['data'].pop('split_ratio')
output_cfg["data"].pop("split_ratio")
comment_dict = {
"device": "Torch device name, e.g., cpu or cuda or cuda:0",
"cpt_path": "Path to the checkpoint file",
"general_pipeline": {"save_path": "Directory to save the inference results"}
"general_pipeline": {
"save_path": "Directory to save the inference results"
},
}
comment_dict = merge_comment(output_cfg, comment_dict)
yaml = ruamel.yaml.YAML()
yaml.dump(comment_dict, Path(cfg).open("w"))
print("Configuration file is generated at {}".format(Path(cfg).absolute()))
print(
"Configuration file is generated at {}".format(
Path(cfg).absolute()
)
)
return config
......@@ -88,8 +108,12 @@ class ApplyNodepredNsPipeline(PipelineBase):
model_name = train_cfg["model_name"]
model_code = NodeModelFactory.get_source_code(model_name)
render_cfg["model_code"] = model_code
render_cfg["model_class_name"] = NodeModelFactory.get_model_class_name(model_name)
render_cfg.update(DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"]))
render_cfg["model_class_name"] = NodeModelFactory.get_model_class_name(
model_name
)
render_cfg.update(
DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"])
)
# Dict for defining cfg in the rendered code
generated_user_cfg = deepcopy(user_cfg_dict)
......
dglgo/dglgo/cli/__init__.py
View file @ dce89919
from .cli import app
if __name__ == "__main__":
app()
\ No newline at end of file
app()
dglgo/dglgo/cli/apply_cli.py
View file @ dce89919
from ..utils.factory import ApplyPipelineFactory
from pathlib import Path
import autopep8
import isort
import typer
import yaml
from pathlib import Path
from ..utils.factory import ApplyPipelineFactory
def apply(
cfg: str = typer.Option(..., help="config yaml file name")
):
def apply(cfg: str = typer.Option(..., help="config yaml file name")):
user_cfg = yaml.safe_load(Path(cfg).open("r"))
pipeline_name = user_cfg["pipeline_name"]
output_file_content = ApplyPipelineFactory.registry[pipeline_name].gen_script(user_cfg)
output_file_content = ApplyPipelineFactory.registry[
pipeline_name
].gen_script(user_cfg)
f_code = autopep8.fix_code(output_file_content, options={'aggressive': 1})
f_code = autopep8.fix_code(output_file_content, options={"aggressive": 1})
f_code = isort.code(f_code)
code = compile(f_code, 'dglgo_tmp.py', 'exec')
exec(code, {'__name__': '__main__'})
code = compile(f_code, "dglgo_tmp.py", "exec")
exec(code, {"__name__": "__main__"})
dglgo/dglgo/cli/cli.py
View file @ dce89919
import typer
from ..pipeline import *
from ..model import *
from .apply_cli import apply
from .config_apply_cli import config_apply_app
from .config_cli import config_app
from .train_cli import train
from .export_cli import export
from .recipe_cli import recipe_app
from .config_apply_cli import config_apply_app
from .apply_cli import apply
from .train_cli import train
no_args_is_help = False
app = typer.Typer(no_args_is_help=True, add_completion=False)
app.add_typer(config_app, name="configure", no_args_is_help=no_args_is_help)
app.add_typer(recipe_app, name="recipe", no_args_is_help=True)
app.command(help="Launch training", no_args_is_help=no_args_is_help)(train)
app.command(help="Export a runnable python script", no_args_is_help=no_args_is_help)(export)
app.add_typer(config_apply_app, name="configure-apply", no_args_is_help=no_args_is_help)
app.command(
help="Export a runnable python script", no_args_is_help=no_args_is_help
)(export)
app.add_typer(
config_apply_app, name="configure-apply", no_args_is_help=no_args_is_help
)
app.command(help="Launch inference", no_args_is_help=no_args_is_help)(apply)
def main():
app()
if __name__ == "__main__":
app()
dglgo/dglgo/cli/config_apply_cli.py
View file @ dce89919
from ..apply_pipeline import *
from ..utils.factory import ApplyPipelineFactory
import typer
config_apply_app = typer.Typer(help="Generate a configuration file for inference")
from ..utils.factory import ApplyPipelineFactory
config_apply_app = typer.Typer(
help="Generate a configuration file for inference"
)
for key, pipeline in ApplyPipelineFactory.registry.items():
config_apply_app.command(key, help=pipeline.get_description())(pipeline.get_cfg_func())
config_apply_app.command(key, help=pipeline.get_description())(
pipeline.get_cfg_func()
)
dglgo/dglgo/cli/config_cli.py
View file @ dce89919
from ..pipeline import *
from ..utils.factory import ModelFactory, PipelineFactory
import typer
from enum import Enum
import typing
import yaml
from enum import Enum
from pathlib import Path
import typer
import yaml
from ..utils.factory import ModelFactory, PipelineFactory
config_app = typer.Typer(help="Generate a configuration file")
for key, pipeline in PipelineFactory.registry.items():
config_app.command(key, help=pipeline.get_description())(pipeline.get_cfg_func())
config_app.command(key, help=pipeline.get_description())(
pipeline.get_cfg_func()
)
if __name__ == "__main__":
config_app()
dglgo/dglgo/cli/export_cli.py
View file @ dce89919
from ..utils.factory import ModelFactory, PipelineFactory, ApplyPipelineFactory
import typer
from enum import Enum
import typing
import yaml
from enum import Enum
from pathlib import Path
import isort
import autopep8
import isort
import typer
import yaml
from ..utils.factory import ApplyPipelineFactory, ModelFactory, PipelineFactory
def export(
cfg: str = typer.Option("cfg.yaml", help="config yaml file name"),
output: str = typer.Option("script.py", help="output python file name")
output: str = typer.Option("script.py", help="output python file name"),
):
user_cfg = yaml.safe_load(Path(cfg).open("r"))
pipeline_name = user_cfg["pipeline_name"]
pipeline_mode = user_cfg["pipeline_mode"]
if pipeline_mode == 'train':
output_file_content = PipelineFactory.registry[pipeline_name].gen_script(user_cfg)
if pipeline_mode == "train":
output_file_content = PipelineFactory.registry[
pipeline_name
].gen_script(user_cfg)
else:
output_file_content = ApplyPipelineFactory.registry[pipeline_name].gen_script(user_cfg)
output_file_content = ApplyPipelineFactory.registry[
pipeline_name
].gen_script(user_cfg)
f_code = autopep8.fix_code(output_file_content, options={'aggressive': 1})
f_code = autopep8.fix_code(output_file_content, options={"aggressive": 1})
f_code = isort.code(f_code)
with open(output, "w") as f:
f.write(f_code)
print("The python script is generated at {}, based on config file {}".format(Path(output).absolute(), Path(cfg).absolute()))
print(
"The python script is generated at {}, based on config file {}".format(
Path(output).absolute(), Path(cfg).absolute()
)
)
if __name__ == "__main__":
export_app = typer.Typer()
export_app.command()(export)
export_app()
\ No newline at end of file
export_app()
dglgo/dglgo/cli/recipe_cli.py
View file @ dce89919
import os
import shutil
from pathlib import Path
from typing import Optional
import typer
import os
import shutil
import yaml
def list_recipes():
file_current_dir = Path(__file__).resolve().parent
recipe_dir = file_current_dir.parent.parent / "recipes"
file_list = list(recipe_dir.glob("*.yaml"))
header = "| {:<30} | {:<18} | {:<20} |".format("Filename", "Pipeline", "Dataset")
typer.echo("="*len(header))
header = "| {:<30} | {:<18} | {:<20} |".format(
"Filename", "Pipeline", "Dataset"
)
typer.echo("=" * len(header))
typer.echo(header)
typer.echo("="*len(header))
typer.echo("=" * len(header))
output_list = []
for file in file_list:
cfg = yaml.safe_load(Path(file).open("r"))
output_list.append({
"file_name": file.name,
"pipeline_name": cfg["pipeline_name"],
"dataset_name": cfg["data"]["name"]
})
output_list.append(
{
"file_name": file.name,
"pipeline_name": cfg["pipeline_name"],
"dataset_name": cfg["data"]["name"],
}
)
# sort by pipeline, if same sort by dataset, if same sort by file name
output_list.sort(key=lambda f: (f["pipeline_name"], f["dataset_name"], f["file_name"]))
output_list.sort(
key=lambda f: (f["pipeline_name"], f["dataset_name"], f["file_name"])
)
for f in output_list:
typer.echo("| {:<30} | {:<18} | {:<20} |".format(f["file_name"], f["pipeline_name"], f["dataset_name"]))
typer.echo("="*len(header))
typer.echo(
"| {:<30} | {:<18} | {:<20} |".format(
f["file_name"], f["pipeline_name"], f["dataset_name"]
)
)
typer.echo("=" * len(header))
def get_recipe(recipe_name: Optional[str] = typer.Argument(None, help="The recipe filename to get, e.q. nodepred_citeseer_gcn.yaml")):
def get_recipe(
recipe_name: Optional[str] = typer.Argument(
None, help="The recipe filename to get, e.q. nodepred_citeseer_gcn.yaml"
)
):
if recipe_name is None:
typer.echo("Usage: dgl recipe get [RECIPE_NAME] \n")
typer.echo(" Copy the recipe to current directory \n")
typer.echo(" Arguments:")
typer.echo(" [RECIPE_NAME] The recipe filename to get, e.q. nodepred_citeseer_gcn.yaml\n")
typer.echo(
" [RECIPE_NAME] The recipe filename to get, e.q. nodepred_citeseer_gcn.yaml\n"
)
typer.echo("Here are all avaliable recipe filename")
list_recipes()
else:
......@@ -41,12 +60,20 @@ def get_recipe(recipe_name: Optional[str] = typer.Argument(None, help="The recip
current_dir = Path(os.getcwd())
recipe_path = recipe_dir / recipe_name
shutil.copy(recipe_path, current_dir)
print("Recipe {} is copied to {}".format(recipe_path.absolute(), current_dir.absolute()))
print(
"Recipe {} is copied to {}".format(
recipe_path.absolute(), current_dir.absolute()
)
)
recipe_app = typer.Typer(help="Get example recipes")
recipe_app.command(name="list", help="List all available example recipes")(list_recipes)
recipe_app.command(name="get", help="Copy the recipe to current directory")(get_recipe)
recipe_app.command(name="list", help="List all available example recipes")(
list_recipes
)
recipe_app.command(name="get", help="Copy the recipe to current directory")(
get_recipe
)
if __name__ == "__main__":
recipe_app()
dglgo/dglgo/cli/train_cli.py
View file @ dce89919
from ..utils.factory import ModelFactory, PipelineFactory
import typer
from enum import Enum
import typing
import yaml
from enum import Enum
from pathlib import Path
import isort
import autopep8
import isort
import typer
import yaml
from ..utils.factory import ModelFactory, PipelineFactory
def train(
cfg: str = typer.Option("cfg.yaml", help="config yaml file name"),
):
user_cfg = yaml.safe_load(Path(cfg).open("r"))
pipeline_name = user_cfg["pipeline_name"]
output_file_content = PipelineFactory.registry[pipeline_name].gen_script(user_cfg)
output_file_content = PipelineFactory.registry[pipeline_name].gen_script(
user_cfg
)
f_code = autopep8.fix_code(output_file_content, options={'aggressive': 1})
f_code = autopep8.fix_code(output_file_content, options={"aggressive": 1})
f_code = isort.code(f_code)
code = compile(f_code, 'dglgo_tmp.py', 'exec')
exec(code, {'__name__': '__main__'})
code = compile(f_code, "dglgo_tmp.py", "exec")
exec(code, {"__name__": "__main__"})
if __name__ == "__main__":
train_app = typer.Typer()
train_app.command()(train)
train_app()
\ No newline at end of file
train_app()
dglgo/dglgo/model/__init__.py
View file @ dce89919
from .node_encoder import *
from .edge_encoder import *
from .graph_encoder import *
\ No newline at end of file
from .graph_encoder import *
dglgo/dglgo/model/edge_encoder/__init__.py
View file @ dce89919
from ...utils.factory import EdgeModelFactory
from .ele import ElementWiseProductPredictor
from .bilinear import BilinearPredictor
from .ele import ElementWiseProductPredictor
EdgeModelFactory.register("ele")(ElementWiseProductPredictor)
EdgeModelFactory.register("bilinear")(BilinearPredictor)
\ No newline at end of file
EdgeModelFactory.register("bilinear")(BilinearPredictor)
dglgo/dglgo/model/edge_encoder/bilinear.py
View file @ dce89919
......@@ -4,11 +4,13 @@ import torch.nn.functional as F
class BilinearPredictor(nn.Module):
def __init__(self,
data_info: dict,
hidden_size: int = 32,
num_layers: int = 1,
bias: bool = True):
def __init__(
self,
data_info: dict,
hidden_size: int = 32,
num_layers: int = 1,
bias: bool = True,
):
"""Bilinear product model for edge scores
Parameters
......@@ -16,7 +18,7 @@ class BilinearPredictor(nn.Module):
data_info : dict
The information about the input dataset.
hidden_size : int
Hidden size.
Hidden size.
num_layers : int
Number of hidden layers.
bias : bool
......@@ -26,7 +28,7 @@ class BilinearPredictor(nn.Module):
in_size, out_size = data_info["in_size"], data_info["out_size"]
self.bilinear = nn.Bilinear(in_size, in_size, hidden_size, bias=bias)
lins_list = []
for _ in range(num_layers-2):
for _ in range(num_layers - 2):
lins_list.append(nn.Linear(hidden_size, hidden_size, bias=bias))
lins_list.append(nn.ReLU())
lins_list.append(nn.Linear(hidden_size, out_size, bias=bias))
......
dglgo/dglgo/model/edge_encoder/dot.py
View file @ dce89919
......@@ -2,16 +2,19 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
class DotPredictor(nn.Module):
def __init__(self,
in_size: int = -1,
out_size: int = 1,
hidden_size: int = 256,
num_layers: int = 3,
bias: bool = False):
def __init__(
self,
in_size: int = -1,
out_size: int = 1,
hidden_size: int = 256,
num_layers: int = 3,
bias: bool = False,
):
super(DotPredictor, self).__init__()
lins_list = []
for _ in range(num_layers-2):
for _ in range(num_layers - 2):
lins_list.append(nn.Linear(in_size, hidden_size, bias=bias))
lins_list.append(nn.ReLU())
lins_list.append(nn.Linear(hidden_size, out_size, bias=bias))
......
dglgo/dglgo/model/edge_encoder/ele.py
View file @ dce89919
......@@ -4,11 +4,13 @@ import torch.nn.functional as F
class ElementWiseProductPredictor(nn.Module):
def __init__(self,
data_info: dict,
hidden_size: int = 64,
num_layers: int = 2,
bias: bool = True):
def __init__(
self,
data_info: dict,
hidden_size: int = 64,
num_layers: int = 2,
bias: bool = True,
):
"""Elementwise product model for edge scores
Parameters
......@@ -16,7 +18,7 @@ class ElementWiseProductPredictor(nn.Module):
data_info : dict
The information about the input dataset.
hidden_size : int
Hidden size.
Hidden size.
num_layers : int
Number of hidden layers.
bias : bool
......
dglgo/dglgo/model/graph_encoder/__init__.py
View file @ dce89919
......@@ -3,4 +3,4 @@ from .gin_ogbg import OGBGGIN
from .pna import PNA
GraphModelFactory.register("gin")(OGBGGIN)
GraphModelFactory.register("pna")(PNA)
\ No newline at end of file
GraphModelFactory.register("pna")(PNA)
dglgo/dglgo/model/graph_encoder/gin_ogbg.py
View file @ dce89919
......@@ -2,12 +2,12 @@ import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import GINEConv, AvgPooling, SumPooling
from dgl.nn import AvgPooling, GINEConv, SumPooling
from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
class MLP(nn.Module):
def __init__(self,
feat_size: int):
def __init__(self, feat_size: int):
"""Multilayer Perceptron (MLP)"""
super(MLP, self).__init__()
self.mlp = nn.Sequential(
......@@ -15,19 +15,22 @@ class MLP(nn.Module):
nn.BatchNorm1d(2 * feat_size),
nn.ReLU(),
nn.Linear(2 * feat_size, feat_size),
nn.BatchNorm1d(feat_size)
nn.BatchNorm1d(feat_size),
)
def forward(self, h):
return self.mlp(h)
class OGBGGIN(nn.Module):
def __init__(self,
data_info: dict,
embed_size: int = 300,
num_layers: int = 5,
dropout: float = 0.5,
virtual_node : bool = False):
def __init__(
self,
data_info: dict,
embed_size: int = 300,
num_layers: int = 5,
dropout: float = 0.5,
virtual_node: bool = False,
):
"""Graph Isomorphism Network (GIN) variant introduced in baselines
for OGB graph property prediction datasets
......@@ -50,21 +53,30 @@ class OGBGGIN(nn.Module):
self.num_layers = num_layers
self.virtual_node = virtual_node
if data_info['name'] in ['ogbg-molhiv', 'ogbg-molpcba']:
if data_info["name"] in ["ogbg-molhiv", "ogbg-molpcba"]:
self.node_encoder = AtomEncoder(embed_size)
self.edge_encoders = nn.ModuleList([
BondEncoder(embed_size) for _ in range(num_layers)])
self.edge_encoders = nn.ModuleList(
[BondEncoder(embed_size) for _ in range(num_layers)]
)
else:
# Handle other datasets
self.node_encoder = nn.Linear(data_info['node_feat_size'], embed_size)
self.edge_encoders = nn.ModuleList([nn.Linear(data_info['edge_feat_size'], embed_size)
for _ in range(num_layers)])
self.conv_layers = nn.ModuleList([GINEConv(MLP(embed_size)) for _ in range(num_layers)])
self.node_encoder = nn.Linear(
data_info["node_feat_size"], embed_size
)
self.edge_encoders = nn.ModuleList(
[
nn.Linear(data_info["edge_feat_size"], embed_size)
for _ in range(num_layers)
]
)
self.conv_layers = nn.ModuleList(
[GINEConv(MLP(embed_size)) for _ in range(num_layers)]
)
self.dropout = nn.Dropout(dropout)
self.pool = AvgPooling()
self.pred = nn.Linear(embed_size, data_info['out_size'])
self.pred = nn.Linear(embed_size, data_info["out_size"])
if virtual_node:
self.virtual_emb = nn.Embedding(1, embed_size)
......
dglgo/dglgo/model/graph_encoder/pna.py
View file @ dce89919
from typing import List
import dgl.function as fn
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import SumPooling, AvgPooling
from dgl.nn import AvgPooling, SumPooling
from ogb.graphproppred.mol_encoder import AtomEncoder
def aggregate_mean(h):
"""mean aggregation"""
return torch.mean(h, dim=1)
def aggregate_max(h):
"""max aggregation"""
return torch.max(h, dim=1)[0]
def aggregate_min(h):
"""min aggregation"""
return torch.min(h, dim=1)[0]
def aggregate_sum(h):
"""sum aggregation"""
return torch.sum(h, dim=1)
def aggregate_var(h):
"""variance aggregation"""
h_mean_squares = torch.mean(h * h, dim=1)
......@@ -30,47 +36,66 @@ def aggregate_var(h):
var = torch.relu(h_mean_squares - h_mean * h_mean)
return var
def aggregate_std(h):
"""standard deviation aggregation"""
return torch.sqrt(aggregate_var(h) + 1e-5)
AGGREGATORS = {'mean': aggregate_mean, 'sum': aggregate_sum, 'max': aggregate_max,
'min': aggregate_min, 'std': aggregate_std, 'var': aggregate_var}
AGGREGATORS = {
"mean": aggregate_mean,
"sum": aggregate_sum,
"max": aggregate_max,
"min": aggregate_min,
"std": aggregate_std,
"var": aggregate_var,
}
def scale_identity(h, D, delta):
"""identity scaling (no scaling operation)"""
return h
def scale_amplification(h, D, delta):
"""amplification scaling"""
return h * (np.log(D + 1) / delta)
def scale_attenuation(h, D, delta):
"""attenuation scaling"""
return h * (delta / np.log(D + 1))
SCALERS = {
'identity': scale_identity,
'amplification': scale_amplification,
'attenuation': scale_attenuation
"identity": scale_identity,
"amplification": scale_amplification,
"attenuation": scale_attenuation,
}
class MLP(nn.Module):
def __init__(self,
in_feat_size: int,
out_feat_size: int,
num_layers: int=3,
decreasing_hidden_size=False):
def __init__(
self,
in_feat_size: int,
out_feat_size: int,
num_layers: int = 3,
decreasing_hidden_size=False,
):
"""Multilayer Perceptron (MLP)"""
super(MLP, self).__init__()
self.layers = nn.ModuleList()
if decreasing_hidden_size:
for i in range(num_layers - 1):
self.layers.append(nn.Linear(in_feat_size // 2 ** i,
in_feat_size // 2 ** (i + 1)))
self.layers.append(nn.Linear(in_feat_size // 2 ** (num_layers - 1),
out_feat_size))
self.layers.append(
nn.Linear(
in_feat_size // 2**i, in_feat_size // 2 ** (i + 1)
)
)
self.layers.append(
nn.Linear(in_feat_size // 2 ** (num_layers - 1), out_feat_size)
)
else:
self.layers.append(nn.Linear(in_feat_size, out_feat_size))
for _ in range(num_layers - 1):
......@@ -84,24 +109,31 @@ class MLP(nn.Module):
h = F.relu(h)
return h
class SimplePNAConv(nn.Module):
r"""A simplified PNAConv variant used in OGB submissions"""
def __init__(self,
feat_size: int,
aggregators: List[str],
scalers: List[str],
delta: float,
dropout: float,
batch_norm: bool,
residual: bool,
num_mlp_layers: int):
def __init__(
self,
feat_size: int,
aggregators: List[str],
scalers: List[str],
delta: float,
dropout: float,
batch_norm: bool,
residual: bool,
num_mlp_layers: int,
):
super(SimplePNAConv, self).__init__()
self.aggregators = [AGGREGATORS[aggr] for aggr in aggregators]
self.scalers = [SCALERS[scale] for scale in scalers]
self.delta = delta
self.mlp = MLP(in_feat_size=(len(aggregators) * len(scalers)) * feat_size,
out_feat_size=feat_size, num_layers=num_mlp_layers)
self.mlp = MLP(
in_feat_size=(len(aggregators) * len(scalers)) * feat_size,
out_feat_size=feat_size,
num_layers=num_mlp_layers,
)
self.dropout = nn.Dropout(dropout)
self.residual = residual
......@@ -111,17 +143,19 @@ class SimplePNAConv(nn.Module):
self.bn = None
def reduce(self, nodes):
h = nodes.mailbox['m']
h = nodes.mailbox["m"]
D = h.shape[-2]
h = torch.cat([aggregate(h) for aggregate in self.aggregators], dim=1)
h = torch.cat([scale(h, D=D, delta=self.delta) for scale in self.scalers], dim=1)
return {'h': h}
h = torch.cat(
[scale(h, D=D, delta=self.delta) for scale in self.scalers], dim=1
)
return {"h": h}
def forward(self, g, h):
with g.local_scope():
g.ndata['h'] = h
g.update_all(fn.copy_u('h', 'm'), self.reduce)
h_new = g.ndata['h']
g.ndata["h"] = h
g.update_all(fn.copy_u("h", "m"), self.reduce)
h_new = g.ndata["h"]
h_new = self.mlp(h_new)
if self.bn is not None:
......@@ -134,18 +168,21 @@ class SimplePNAConv(nn.Module):
return h_new
class PNA(nn.Module):
def __init__(self,
data_info: dict,
embed_size: int = 80,
aggregators: str = 'mean max min std',
scalers: str = 'identity amplification attenuation',
dropout: float = 0.3,
batch_norm: bool = True,
residual: bool = True,
num_mlp_layers: int = 1,
num_layers: int = 4,
readout: str = 'mean'):
def __init__(
self,
data_info: dict,
embed_size: int = 80,
aggregators: str = "mean max min std",
scalers: str = "identity amplification attenuation",
dropout: float = 0.3,
batch_norm: bool = True,
residual: bool = True,
num_mlp_layers: int = 1,
num_layers: int = 4,
readout: str = "mean",
):
"""Principal Neighbourhood Aggregation
Parameters
......@@ -182,44 +219,62 @@ class PNA(nn.Module):
self.readout = readout
if aggregators is None:
aggregators = ['mean', 'max', 'min', 'std']
aggregators = ["mean", "max", "min", "std"]
else:
aggregators = [agg.strip() for agg in aggregators.split(' ')]
assert set(aggregators).issubset({'mean', 'max', 'min', 'std', 'sum'}), \
"Expect aggregators to be a subset of ['mean', 'max', 'min', 'std', 'sum'], \
got {}".format(aggregators)
aggregators = [agg.strip() for agg in aggregators.split(" ")]
assert set(aggregators).issubset(
{"mean", "max", "min", "std", "sum"}
), "Expect aggregators to be a subset of ['mean', 'max', 'min', 'std', 'sum'], \
got {}".format(
aggregators
)
if scalers is None:
scalers = ['identity', 'amplification', 'attenuation']
scalers = ["identity", "amplification", "attenuation"]
else:
scalers = [scl.strip() for scl in scalers.split(' ')]
assert set(scalers).issubset({'identity', 'amplification', 'attenuation'}), \
"Expect scalers to be a subset of ['identity', 'amplification', 'attenuation'], \
got {}".format(scalers)
scalers = [scl.strip() for scl in scalers.split(" ")]
assert set(scalers).issubset(
{"identity", "amplification", "attenuation"}
), "Expect scalers to be a subset of ['identity', 'amplification', 'attenuation'], \
got {}".format(
scalers
)
self.aggregators = aggregators
self.scalers = scalers
if data_info['name'] in ['ogbg-molhiv', 'ogbg-molpcba']:
if data_info["name"] in ["ogbg-molhiv", "ogbg-molpcba"]:
self.node_encoder = AtomEncoder(embed_size)
else:
# Handle other datasets
self.node_encoder = nn.Linear(data_info['node_feat_size'], embed_size)
self.conv_layers = nn.ModuleList([SimplePNAConv(feat_size=embed_size,
aggregators=aggregators,
scalers=scalers,
delta=data_info['delta'],
dropout=dropout,
batch_norm=batch_norm,
residual=residual,
num_mlp_layers=num_mlp_layers)
for _ in range(num_layers)])
if readout == 'sum':
self.node_encoder = nn.Linear(
data_info["node_feat_size"], embed_size
)
self.conv_layers = nn.ModuleList(
[
SimplePNAConv(
feat_size=embed_size,
aggregators=aggregators,
scalers=scalers,
delta=data_info["delta"],
dropout=dropout,
batch_norm=batch_norm,
residual=residual,
num_mlp_layers=num_mlp_layers,
)
for _ in range(num_layers)
]
)
if readout == "sum":
self.pool = SumPooling()
elif readout == 'mean':
elif readout == "mean":
self.pool = AvgPooling()
else:
raise ValueError("Expect readout to be 'sum' or 'mean', got {}".format(readout))
self.pred = MLP(embed_size, data_info['out_size'], decreasing_hidden_size=True)
raise ValueError(
"Expect readout to be 'sum' or 'mean', got {}".format(readout)
)
self.pred = MLP(
embed_size, data_info["out_size"], decreasing_hidden_size=True
)
def forward(self, graph, node_feat, edge_feat=None):
hn = self.node_encoder(node_feat)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment