Merge pull request #466 from nqhq-lou/dev_wandb

Add Command-Line Argument for Specifying Wandb Directory

mace/cli/run_train.py

@@ -688,6 +688,7 @@ def run(args: argparse.Namespace) -> None:
             entity=args.wandb_entity,
             name=args.wandb_name,
             config=wandb_config,
+            directory=args.wandb_dir,
         )
         wandb.run.summary["params"] = args_dict_json
 

mace/tools/arg_parser.py

@@ -574,6 +574,12 @@ def build_default_arg_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--wandb_dir",
+        help="An absolute path to a directory where Weights and Biases metadata will be stored",
+        type=str,
+        default=None,
+    )
     parser.add_argument(
         "--wandb_project",
         help="Weights and Biases project name",

mace/tools/torch_tools.py

@@ -119,10 +119,10 @@ def voigt_to_matrix(t: torch.Tensor):
     )
 
 
-def init_wandb(project: str, entity: str, name: str, config: dict):
+def init_wandb(project: str, entity: str, name: str, config: dict, directory:str):
     import wandb
 
-    wandb.init(project=project, entity=entity, name=name, config=config)
+    wandb.init(project=project, entity=entity, name=name, config=config, dir=directory)
 
 
 @contextmanager