Simplify examples

examples/04_training/01_train_model.py

@@ -19,11 +19,6 @@
 from graphnet.utilities.logging import Logger
 
 
-# Make sure W&B output directory exists
-WANDB_DIR = "./wandb/"
-os.makedirs(WANDB_DIR, exist_ok=True)
-
-
 def main(
     dataset_config_path: str,
     model_config_path: str,
@@ -34,18 +29,23 @@ def main(
     num_workers: int,
     prediction_names: Optional[List[str]],
     suffix: Optional[str] = None,
+    wandb: bool = False,
 ) -> None:
     """Run example."""
     # Construct Logger
     logger = Logger()
 
     # Initialise Weights & Biases (W&B) run
-    wandb_logger = WandbLogger(
-        project="example-script",
-        entity="graphnet-team",
-        save_dir=WANDB_DIR,
-        log_model=True,
-    )
+    if wandb:
+        # Make sure W&B output directory exists
+        wandb_dir = "./wandb/"
+        os.makedirs(wandb_dir, exist_ok=True)
+        wandb_logger = WandbLogger(
+            project="example-script",
+            entity="graphnet-team",
+            save_dir=wandb_dir,
+            log_model=True,
+        )
 
     # Build model
     model_config = ModelConfig.load(model_config_path)
@@ -80,7 +80,7 @@ def main(
     # Log configurations to W&B
     # NB: Only log to W&B on the rank-zero process in case of multi-GPU
     #     training.
-    if rank_zero_only.rank == 0:
+    if wandb and rank_zero_only.rank == 0:
         wandb_logger.experiment.config.update(config)
         wandb_logger.experiment.config.update(model_config.as_dict())
         wandb_logger.experiment.config.update(dataset_config.as_dict())
@@ -98,7 +98,7 @@ def main(
         dataloaders["train"],
         dataloaders["validation"],
         callbacks=callbacks,
-        logger=wandb_logger,
+        logger=wandb_logger if wandb else None,
         **config.fit,
     )
 
@@ -166,6 +166,12 @@ def main(
         default=None,
     )
 
+    parser.add_argument(
+        "--wandb",
+        action="store_true",
+        help="If True, Weights & Biases are used to track the experiment.",
+    )
+
     args = parser.parse_args()
 
     main(
@@ -178,4 +184,5 @@ def main(
         args.num_workers,
         args.prediction_names,
         args.suffix,
+        args.wandb,
     )

examples/04_training/02_train_model_without_configs.py

@@ -25,10 +25,6 @@
 features = FEATURES.PROMETHEUS
 truth = TRUTH.PROMETHEUS
 
-# Make sure W&B output directory exists
-WANDB_DIR = "./wandb/"
-os.makedirs(WANDB_DIR, exist_ok=True)
-
 
 def main(
     path: str,
@@ -40,18 +36,23 @@ def main(
     early_stopping_patience: int,
     batch_size: int,
     num_workers: int,
+    wandb: bool = False,
 ) -> None:
     """Run example."""
     # Construct Logger
     logger = Logger()
 
     # Initialise Weights & Biases (W&B) run
-    wandb_logger = WandbLogger(
-        project="example-script",
-        entity="graphnet-team",
-        save_dir=WANDB_DIR,
-        log_model=True,
-    )
+    if wandb:
+        # Make sure W&B output directory exists
+        wandb_dir = "./wandb/"
+        os.makedirs(wandb_dir, exist_ok=True)
+        wandb_logger = WandbLogger(
+            project="example-script",
+            entity="graphnet-team",
+            save_dir=wandb_dir,
+            log_model=True,
+        )
 
     logger.info(f"features: {features}")
     logger.info(f"truth: {truth}")
@@ -72,9 +73,9 @@ def main(
 
     archive = os.path.join(EXAMPLE_OUTPUT_DIR, "train_model_without_configs")
     run_name = "dynedge_{}_example".format(config["target"])
-
-    # Log configuration to W&B
-    wandb_logger.experiment.config.update(config)
+    if wandb:
+        # Log configuration to W&B
+        wandb_logger.experiment.config.update(config)
 
     (
         training_dataloader,
@@ -137,17 +138,16 @@ def main(
         training_dataloader,
         validation_dataloader,
         callbacks=callbacks,
-        logger=wandb_logger,
+        logger=wandb_logger if wandb else None,
         **config["fit"],
     )
 
     # Get predictions
-    prediction_columns = [config["target"] + "_pred"]
-    additional_attributes = [config["target"]]
+    additional_attributes = model.target_labels
+    assert isinstance(additional_attributes, list)  # mypy
 
     results = model.predict_as_dataframe(
         validation_dataloader,
-        prediction_columns=prediction_columns,
         additional_attributes=additional_attributes + ["event_no"],
     )
 
@@ -206,6 +206,12 @@ def main(
         "num-workers",
     )
 
+    parser.add_argument(
+        "--wandb",
+        action="store_true",
+        help="If True, Weights & Biases are used to track the experiment.",
+    )
+
     args = parser.parse_args()
 
     main(
@@ -218,4 +224,5 @@ def main(
         args.early_stopping_patience,
         args.batch_size,
         args.num_workers,
+        args.wandb,
     )

src/graphnet/models/model.py

@@ -10,6 +10,7 @@
 import pandas as pd
 from pytorch_lightning import Trainer, LightningModule
 from pytorch_lightning.callbacks.callback import Callback
+from pytorch_lightning.callbacks import EarlyStopping
 from pytorch_lightning.loggers.logger import Logger as LightningLogger
 import torch
 from torch import Tensor
@@ -18,6 +19,7 @@
 
 from graphnet.utilities.logging import Logger
 from graphnet.utilities.config import Configurable, ModelConfig
+from graphnet.training.callbacks import ProgressBar
 
 
 class Model(Logger, Configurable, LightningModule, ABC):
@@ -85,9 +87,17 @@ def fit(
         log_every_n_steps: int = 1,
         gradient_clip_val: Optional[float] = None,
         distribution_strategy: Optional[str] = "ddp",
+        early_stopping_patience: int = 5,
         **trainer_kwargs: Any,
     ) -> None:
         """Fit `Model` using `pytorch_lightning.Trainer`."""
+        # Checks
+        if callbacks is None:
+            callbacks = self._create_default_callbacks(
+                val_dataloader=val_dataloader,
+                early_stopping_patience=early_stopping_patience,
+            )
+
         self.train(mode=True)
 
         self._construct_trainers(
@@ -110,6 +120,26 @@ def fit(
             self.warning("[ctrl+c] Exiting gracefully.")
             pass
 
+    def _create_default_callbacks(
+        self, val_dataloader: DataLoader, early_stopping_patience: int
+    ) -> List:
+        callbacks = [ProgressBar()]
+        if val_dataloader is not None:
+            has_es = False
+            assert isinstance(callbacks, list)
+            for callback in callbacks:
+                if isinstance(callback, EarlyStopping):
+                    has_es = True
+            if has_es is False:
+                callbacks.append(
+                    EarlyStopping(
+                        monitor="val_loss",
+                        patience=early_stopping_patience,
+                    )
+                )
+                self.info("EarlyStopping callback added automatically.")
+        return callbacks
+
     def predict(
         self,
         dataloader: DataLoader,
@@ -146,7 +176,7 @@ def predict(
     def predict_as_dataframe(
         self,
         dataloader: DataLoader,
-        prediction_columns: List[str],
+        prediction_columns: Optional[List[str]] = None,
         *,
         node_level: bool = False,
         additional_attributes: Optional[List[str]] = None,
@@ -164,6 +194,14 @@ def predict_as_dataframe(
             additional_attributes = []
         assert isinstance(additional_attributes, list)
 
+        if prediction_columns is None:
+            try:
+                prediction_columns = self.prediction_columns
+            except AttributeError:
+                assert (
+                    1 == 2
+                ), "Could not infer prediction_columns from model. Please specify prediction_columns."
+        assert isinstance(prediction_columns, list)
         if (
             not isinstance(dataloader.sampler, SequentialSampler)
             and additional_attributes
@@ -178,6 +216,7 @@ def predict_as_dataframe(
                 "doesn't resample batches; or do not request "
                 "`additional_attributes`."
             )
+        self.info(f"Column names for predictions are: \n {prediction_columns}")
         predictions_torch = self.predict(
             dataloader=dataloader,
             gpus=gpus,

src/graphnet/models/standard_model.py

@@ -62,6 +62,16 @@ def __init__(
         self._scheduler_kwargs = scheduler_kwargs or dict()
         self._scheduler_config = scheduler_config or dict()
 
+    @property
+    def target_labels(self) -> List[str]:
+        """Return target label."""
+        return [label for task in self._tasks for label in task._target_labels]
+
+    @property
+    def prediction_labels(self) -> List[str]:
+        """Return prediction labels."""
+        return [label for task in self._tasks for label in task._output_labels]
+
     def configure_optimizers(self) -> Dict[str, Any]:
         """Configure the model's optimizer(s)."""
         optimizer = self._optimizer_class(