Add --no_log and --force_log options.

--no_log can be used for debug runs that you don't want appearing on W&B, e.g. if there
may be errors in the script that happen after the W&B run starts and you want to fix
those first, or if reviewing a PR and testing the training to make sure it runs.

--force_log can be used during `model test` to create W&B run so the metrics can be
viewed there, the primary use case for now is if you log a confusion matrix, if it wasn't
there during training or if val/test are different splits then it needs to be in W&B to
visualize easily.

Author: uakfdotb

rslp/lightning_cli.py

@@ -77,6 +77,18 @@ def add_arguments_to_parser(self, parser) -> None:
             help="Load best checkpoint from GCS for test/predict",
             default=False,
         )
+        parser.add_argument(
+            "--force_log",
+            type=bool,
+            help="Log to W&B even for test/predict",
+            default=False,
+        )
+        parser.add_argument(
+            "--no_log",
+            type=bool,
+            help="Disable W&B logging for fit",
+            default=False,
+        )
 
     def before_instantiate_classes(self):
         """Called before Lightning class initialization."""
@@ -92,7 +104,7 @@ def before_instantiate_classes(self):
             )
         )
 
-        if subcommand == "fit":
+        if (subcommand == "fit" and not c.no_log) or c.force_log:
             # Add and configure WandbLogger as needed.
             if not c.trainer.logger:
                 c.trainer.logger = jsonargparse.Namespace(
@@ -114,6 +126,7 @@ def before_instantiate_classes(self):
                 }
             )
 
+        if subcommand == "fit" and not c.no_log:
             # Set the checkpoint directory to canonical GCS location.
             checkpoint_callback = None
             upload_wandb_callback = None