[DGL-Go] Change name to dglgo (#3778)

* add

* remove

* fix

* rework the readme and some changes

* add png

* update png

* add recipe get
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>

enter/dglenter/pipeline/linkpred/gen.py → dglgo/dglgo/pipeline/linkpred/gen.py

@@ -20,6 +20,7 @@ class LinkpredPipelineCfg(BaseModel):
     eval_period: int = 5
     optimizer: dict = {"name": "Adam", "lr": 0.005}
     loss: str = "BCELoss"
+    save_path: str = "model.pth"
     num_runs: int = 1
 
 
@@ -29,6 +30,7 @@ pipeline_comments = {
     "train_batch_size": "Edge batch size when training",
     "num_epochs": "Number of training epochs",
     "eval_period": "Interval epochs between evaluations",
+    "save_path": "Path to save the model",
     "num_runs": "Number of experiments to run",
 }
 
@@ -67,20 +69,18 @@ class LinkpredPipeline(PipelineBase):
         def config(
             data: DataFactory.filter("linkpred").get_dataset_enum() = typer.Option(..., help="input data name"),
             cfg: str = typer.Option(
-                "cfg.yml", help="output configuration path"),
+                "cfg.yaml", help="output configuration path"),
             node_model: NodeModelFactory.get_model_enum() = typer.Option(...,
                                                                          help="Model name"),
             edge_model: EdgeModelFactory.get_model_enum() = typer.Option(...,
                                                                          help="Model name"),
             neg_sampler: NegativeSamplerFactory.get_model_enum() = typer.Option(
-                "uniform", help="Negative sampler name"),
-            device: DeviceEnum = typer.Option(
-                "cpu", help="Device, cpu or cuda"),
+                "persource", help="Negative sampler name"),
         ):
             self.__class__.setup_user_cfg_cls()
             generated_cfg = {
                 "pipeline_name": "linkpred",
-                "device": device.value,
+                "device": "cpu",
                 "data": {"name": data.name},
                 "neg_sampler": {"name": neg_sampler.value},
                 "node_model": {"name": node_model.value},
@@ -89,6 +89,7 @@ class LinkpredPipeline(PipelineBase):
             output_cfg = self.user_cfg_cls(**generated_cfg).dict()
             output_cfg = deep_convert_dict(output_cfg)
             comment_dict = {
+                "device": "Torch device name, e.q. cpu or cuda or cuda:0",
                 "general_pipeline": pipeline_comments,
                 "node_model": NodeModelFactory.get_constructor_doc_dict(node_model.value),
                 "edge_model": EdgeModelFactory.get_constructor_doc_dict(edge_model.value),
@@ -99,6 +100,9 @@ class LinkpredPipeline(PipelineBase):
                 },
             }
             comment_dict = merge_comment(output_cfg, comment_dict)
+
+            if cfg is None:
+                cfg = "_".join(["linkpred", data.value, node_model.value, edge_model.value]) + ".yaml"
             yaml = ruamel.yaml.YAML()
             yaml.dump(comment_dict, Path(cfg).open("w"))
             print("Configuration file is generated at {}".format(Path(cfg).absolute()))

enter/dglenter/pipeline/linkpred/linkpred.jinja-py → dglgo/dglgo/pipeline/linkpred/linkpred.jinja-py

@@ -112,6 +112,7 @@ def main():
     loss = torch.nn.{{ loss }}()
     optimizer = torch.optim.Adam(params, **pipeline_cfg["optimizer"])
     test_hits = train(cfg, pipeline_cfg, device, dataset, model, optimizer, loss)
+    torch.save(model, pipeline_cfg["save_path"])
     return test_hits
 
 if __name__ == '__main__':

enter/dglenter/pipeline/nodepred/gen.py → dglgo/dglgo/pipeline/nodepred/gen.py

@@ -18,6 +18,7 @@ pipeline_comments = {
         "patience": "Steps before early stop",
         "checkpoint_path": "Early stop checkpoint model file path"
     },
+    "save_path": "Path to save the model",
     "num_runs": "Number of experiments to run",
 }
 
@@ -27,6 +28,7 @@ class NodepredPipelineCfg(BaseModel):
     eval_period: int = 5
     optimizer: dict = {"name": "Adam", "lr": 0.01, "weight_decay": 5e-4}
     loss: str = "CrossEntropyLoss"
+    save_path: str = "model.pth"
     num_runs: int = 1
 
 @PipelineFactory.register("nodepred")
@@ -54,15 +56,14 @@ class NodepredPipeline(PipelineBase):
     def get_cfg_func(self):
         def config(
             data: DataFactory.filter("nodepred").get_dataset_enum() = typer.Option(..., help="input data name"),
-            cfg: str = typer.Option(
-                "cfg.yml", help="output configuration path"),
+            cfg: Optional[str] = typer.Option(
+                None, help="output configuration path"),
             model: NodeModelFactory.get_model_enum() = typer.Option(..., help="Model name"),
-            device: DeviceEnum = typer.Option("cpu", help="Device, cpu or cuda"),
         ):  
             self.__class__.setup_user_cfg_cls()
             generated_cfg = {
                 "pipeline_name": self.pipeline_name,
-                "device": device,
+                "device": "cpu",
                 "data": {"name": data.name},
                 "model": {"name": model.value},
                 "general_pipeline": {}
@@ -70,6 +71,7 @@ class NodepredPipeline(PipelineBase):
             output_cfg = self.user_cfg_cls(**generated_cfg).dict()
             output_cfg = deep_convert_dict(output_cfg)
             comment_dict = {
+                "device": "Torch device name, e.q. cpu or cuda or cuda:0",
                 "data": {
                     "split_ratio": 'Ratio to generate split masks, for example set to [0.8, 0.1, 0.1] for 80% train/10% val/10% test. Leave blank to use builtin split in original dataset'
                 },
@@ -79,6 +81,8 @@ class NodepredPipeline(PipelineBase):
             comment_dict = merge_comment(output_cfg, comment_dict)
 
             yaml = ruamel.yaml.YAML()
+            if cfg is None:
+                cfg = "_".join(["nodepred", data.value, model.value]) + ".yaml"
             yaml.dump(comment_dict, Path(cfg).open("w"))
             print("Configuration file is generated at {}".format(Path(cfg).absolute()))
 
@@ -88,7 +92,7 @@ class NodepredPipeline(PipelineBase):
     def gen_script(cls, user_cfg_dict):
         # Check validation
         cls.setup_user_cfg_cls()
-        user_cfg = cls.user_cfg_cls(**user_cfg_dict)
+        user_cfg = cls.user_cfg_cls(**user_cfg_dict)        
         file_current_dir = Path(__file__).resolve().parent
         with open(file_current_dir / "nodepred.jinja-py", "r") as f:
             template = Template(f.read())
@@ -102,6 +106,8 @@ class NodepredPipeline(PipelineBase):
         render_cfg.update(DataFactory.get_generated_code_dict(user_cfg_dict["data"]["name"], '**cfg["data"]'))
 
         generated_user_cfg = copy.deepcopy(user_cfg_dict)
+        if "split_ratio" in generated_user_cfg["data"]:
+            generated_user_cfg["data"].pop("split_ratio")
         if len(generated_user_cfg["data"]) == 1:
             generated_user_cfg.pop("data")
         else:
@@ -116,9 +122,6 @@ class NodepredPipeline(PipelineBase):
 
         if user_cfg_dict["data"].get("split_ratio", None) is not None:
             render_cfg["data_initialize_code"] = "{}, split_ratio={}".format(render_cfg["data_initialize_code"], user_cfg_dict["data"]["split_ratio"])
-        if "split_ratio" in generated_user_cfg["data"]:
-            generated_user_cfg["data"].pop("split_ratio")
-
         render_cfg["user_cfg_str"] = f"cfg = {str(generated_user_cfg)}"
         render_cfg["user_cfg"] = user_cfg_dict
         return template.render(**render_cfg)

enter/dglenter/pipeline/nodepred/nodepred.jinja-py → dglgo/dglgo/pipeline/nodepred/nodepred.jinja-py

@@ -112,6 +112,7 @@ def main():
     optimizer = torch.optim.{{ user_cfg.general_pipeline.optimizer.name }}(model.parameters(), **pipeline_cfg["optimizer"])
     # train
     test_acc = train(cfg, pipeline_cfg, device, data, model, optimizer, loss)
+    torch.save(model, pipeline_cfg["save_path"])
     return test_acc
 
 if __name__ == '__main__':

enter/dglenter/pipeline/nodepred_sample/gen.py → dglgo/dglgo/pipeline/nodepred_sample/gen.py

@@ -36,6 +36,14 @@ pipeline_comments = {
         "patience": "Steps before early stop",
         "checkpoint_path": "Early stop checkpoint model file path"
     },
+    "sampler": {
+        "fan_out": "List of neighbors to sample per edge type for each GNN layer, with the i-th element being the fanout for the i-th GNN layer. Length should be the same as num_layers in model setting",
+        "batch_size": "Batch size of seed nodes in training stage",
+        "num_workers": "Number of workers to accelerate the graph data processing step",
+        "eval_batch_size": "Batch size of seed nodes in training stage in evaluation stage",
+        "eval_num_workers": "Number of workers to accelerate the graph data processing step in evaluation stage"
+    },
+    "save_path": "Path to save the model",
     "num_runs": "Number of experiments to run",
 }
 
@@ -47,6 +55,7 @@ class NodepredNSPipelineCfg(BaseModel):
     optimizer: dict = {"name": "Adam", "lr": 0.005, "weight_decay": 0.0}
     loss: str = "CrossEntropyLoss"
     num_runs: int = 1
+    save_path: str = "model.pth"
 
 @PipelineFactory.register("nodepred-ns")
 class NodepredNsPipeline(PipelineBase):
@@ -60,7 +69,7 @@ class NodepredNsPipeline(PipelineBase):
         class NodePredUserConfig(UserConfig):
             eval_device: DeviceEnum = Field("cpu")
             data: DataFactory.filter("nodepred-ns").get_pydantic_config() = Field(..., discriminator="name")
-            model : NodeModelFactory.get_pydantic_model_config() = Field(..., discriminator="name")   
+            model : NodeModelFactory.filter(lambda cls: hasattr(cls, "forward_block")).get_pydantic_model_config() = Field(..., discriminator="name")   
             general_pipeline: NodepredNSPipelineCfg
 
         cls.user_cfg_cls = NodePredUserConfig
@@ -72,16 +81,14 @@ class NodepredNsPipeline(PipelineBase):
     def get_cfg_func(self):
         def config(
             data: DataFactory.filter("nodepred-ns").get_dataset_enum() = typer.Option(..., help="input data name"),
-            cfg: str = typer.Option(
-                "cfg.yml", help="output configuration path"),
-            model: NodeModelFactory.get_model_enum() = typer.Option(..., help="Model name"),
-            device: DeviceEnum = typer.Option(
-                "cpu", help="Device, cpu or cuda"),
+            cfg: Optional[str] = typer.Option(
+                None, help="output configuration path"),
+            model: NodeModelFactory.filter(lambda cls: hasattr(cls, "forward_block")).get_model_enum() = typer.Option(..., help="Model name"),
         ):
             self.__class__.setup_user_cfg_cls()
-            generated_cfg = {
+            generated_cfg = {                
                 "pipeline_name": "nodepred-ns",
-                "device": device,
+                "device": "cpu",
                 "data": {"name": data.name},
                 "model": {"name": model.value},
                 "general_pipeline": {"sampler":{"name": "neighbor"}}
@@ -89,14 +96,21 @@ class NodepredNsPipeline(PipelineBase):
             output_cfg = self.user_cfg_cls(**generated_cfg).dict()
             output_cfg = deep_convert_dict(output_cfg)
             comment_dict = {
+                "device": "Torch device name, e.q. cpu or cuda or cuda:0",
                 "data": {
                     "split_ratio": 'Ratio to generate split masks, for example set to [0.8, 0.1, 0.1] for 80% train/10% val/10% test. Leave blank to use builtin split in original dataset'
                 },
                 "general_pipeline": pipeline_comments,
-                "model": NodeModelFactory.get_constructor_doc_dict(model.value)
+                "model": NodeModelFactory.get_constructor_doc_dict(model.value),
             }
             comment_dict = merge_comment(output_cfg, comment_dict)
 
+            # truncate length fan_out to be the same as num_layers in model
+            if "num_layers" in comment_dict["model"]:
+                comment_dict['general_pipeline']["sampler"]["fan_out"] = [5,10,15,15,15][:int(comment_dict['model']["num_layers"])]
+
+            if cfg is None:
+                cfg = "_".join(["nodepred-ns", data.value, model.value]) + ".yaml"
             yaml = ruamel.yaml.YAML()
             yaml.dump(comment_dict, Path(cfg).open("w"))
             print("Configuration file is generated at {}".format(
@@ -112,6 +126,10 @@ class NodepredNsPipeline(PipelineBase):
             template = Template(f.read())
         pipeline_cfg = NodepredNSPipelineCfg(
             **user_cfg_dict["general_pipeline"])
+        
+        if "num_layers" in user_cfg_dict["model"]:
+            assert user_cfg_dict["model"]["num_layers"] == len(user_cfg_dict["general_pipeline"]["sampler"]["fan_out"]), \
+                "The num_layers in model config should be the same as the length of fan_out in sampler. For example, if num_layers is 1, the fan_out cannot be [5, 10]"          
 
         render_cfg = copy.deepcopy(user_cfg_dict)
         model_code = NodeModelFactory.get_source_code(
@@ -123,6 +141,8 @@ class NodepredNsPipeline(PipelineBase):
             user_cfg_dict["data"]["name"], '**cfg["data"]'))
         generated_user_cfg = copy.deepcopy(user_cfg_dict)
 
+        if "split_ratio" in generated_user_cfg["data"]:
+            generated_user_cfg["data"].pop("split_ratio")
         if len(generated_user_cfg["data"]) == 1:
             generated_user_cfg.pop("data")
         else:
@@ -135,8 +155,6 @@ class NodepredNsPipeline(PipelineBase):
 
         if user_cfg_dict["data"].get("split_ratio", None) is not None:
             render_cfg["data_initialize_code"] = "{}, split_ratio={}".format(render_cfg["data_initialize_code"], user_cfg_dict["data"]["split_ratio"])
-        if "split_ratio" in generated_user_cfg["data"]:
-            generated_user_cfg["data"].pop("split_ratio")
 
         render_cfg["user_cfg_str"] = f"cfg = {str(generated_user_cfg)}"
         render_cfg["user_cfg"] = user_cfg_dict
@@ -145,4 +163,4 @@ class NodepredNsPipeline(PipelineBase):
 
     @staticmethod
     def get_description() -> str:
-        return "Node classification sampling pipeline"
+        return "Node classification neighbor sampling pipeline"

enter/dglenter/pipeline/nodepred_sample/nodepred-ns.jinja-py → dglgo/dglgo/pipeline/nodepred_sample/nodepred-ns.jinja-py

@@ -157,8 +157,8 @@ def main():
     model = model.to(device)
     loss = torch.nn.{{ user_cfg.general_pipeline.loss }}()
     optimizer = torch.optim.{{ user_cfg.general_pipeline.optimizer.name }}(model.parameters(), **pipeline_cfg["optimizer"])
-    # train
     test_acc = train(cfg, pipeline_cfg, device, data, model, optimizer, loss)
+    torch.save(model, pipeline_cfg["save_path"])
     return test_acc
 
 if __name__ == '__main__':

enter/dglenter/utils/factory.py → dglgo/dglgo/utils/factory.py

@@ -334,6 +334,14 @@ class ModelFactory:
             type_annotation_dict[k] = param.annotation
         return type_annotation_dict
 
+    def filter(self, filter_func):
+        new_fac = ModelFactory()
+        for name in self.registry:
+            if filter_func(self.registry[name]):
+                new_fac.registry[name] = self.registry[name]
+                new_fac.code_registry[name] = self.code_registry[name]
+        return new_fac
+
 
 class SamplerFactory:
     """ The factory class for creating executors"""
@@ -411,7 +419,7 @@ class SamplerFactory:
 
 
 NegativeSamplerFactory = SamplerFactory()
-NegativeSamplerFactory.register("uniform")(GlobalUniform)
+NegativeSamplerFactory.register("global")(GlobalUniform)
 NegativeSamplerFactory.register("persource")(PerSourceUniform)
 
 NodeModelFactory = ModelFactory()

enter/recipes/linkpred_citation2_sage.yml → dglgo/recipes/linkpred_citation2_sage.yaml

@@ -31,4 +31,5 @@ general_pipeline:
     name: Adam
     lr: 0.005
   loss: BCELoss
+  save_path: "model.pth"
   num_runs: 1                 # Number of experiments to run

enter/recipes/linkpred_collab_sage.yml → dglgo/recipes/linkpred_collab_sage.yaml

@@ -31,4 +31,5 @@ general_pipeline:
     name: Adam
     lr: 0.005
   loss: BCELoss
+  save_path: "model.pth"
   num_runs: 1                 # Number of experiments to run