feat(wandb): log models as artifacts

pytorch_lightning/loggers/base.py

@@ -26,6 +26,7 @@
 
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning import Trainer
 
 
 def rank_zero_experiment(fn: Callable) -> Callable:
@@ -71,6 +72,15 @@ def __init__(
         self._agg_key_funcs = agg_key_funcs if agg_key_funcs else {}
         self._agg_default_func = agg_default_func
 
+    def connect(self, trainer: Optional[Trainer] = None) -> None:
+        """
+        Connect trainer to logger
+
+        Args:
+            trainer: lightning Trainer
+        """
+        pass
+
     def update_agg_funcs(
         self,
         agg_key_funcs: Optional[Mapping[str, Callable[[Sequence[float]], float]]] = None,
@@ -355,6 +365,10 @@ def __init__(self, logger_iterable: Iterable[LightningLoggerBase]):
     def __getitem__(self, index: int) -> LightningLoggerBase:
         return [logger for logger in self._logger_iterable][index]
 
+    def connect(self, trainer: Optional[Trainer] = None) -> None:
+        for logger in self._logger_iterable:
+            logger.connect(trainer)
+
     def update_agg_funcs(
         self,
         agg_key_funcs: Optional[Mapping[str, Callable[[Sequence[float]], float]]] = None,

pytorch_lightning/loggers/wandb.py

@@ -16,6 +16,8 @@
 -------------------------
 """
 import os
+import re
+import numbers
 from argparse import Namespace
 from typing import Any, Dict, Optional, Union
 
@@ -168,10 +170,6 @@ def experiment(self) -> Run:
                 **self._kwargs
             ) if wandb.run is None else wandb.run
 
-            # save checkpoints in wandb dir to upload on W&B servers
-            if self._save_dir is None:
-                self._save_dir = self._experiment.dir
-
             # define default x-axis (for latest wandb versions)
             if getattr(self._experiment, "define_metric", None):
                 self._experiment.define_metric("trainer/global_step")
@@ -215,6 +213,18 @@ def version(self) -> Optional[str]:
 
     @rank_zero_only
     def finalize(self, status: str) -> None:
-        # upload all checkpoints from saving dir
+        # save checkpoints as artifacts
         if self._log_model:
-            wandb.save(os.path.join(self.save_dir, "*.ckpt"))
+            # use run name and ensure it's a valid Artifact name
+            artifact_name = re.sub(r"[^a-zA-Z0-9_\.\-]", "", self.experiment.name)
+            # gather interesting metadata
+            metadata = {
+                k: v
+                for k, v in dict(self.experiment.summary).items()
+                if isinstance(v, numbers.Number) and not k.startswith("_")
+            }
+            # TODO: see if we can also log data from `trainer.checkpoint_callback` (best_model_path, etc) 
+            artifact = wandb.Artifact(name=f"run-{artifact_name}", type="model", metadata=metadata)
+            # TODO: we need access to `trainer.checkpoint_callback.dirpath`
+            artifact.add_dir(trainer.checkpoint_callback.dirpath)
+            self.experiment.log_artifact(artifact)

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -143,6 +143,7 @@ def should_update_logs(self):
         return should_log_every_n_steps or self.trainer.should_stop
 
     def configure_logger(self, logger):
+        # connect logger to trainer
         if logger is True:
             version = os.environ.get('PL_EXP_VERSION', self.trainer.slurm_job_id)
 
@@ -158,6 +159,9 @@ def configure_logger(self, logger):
             else:
                 self.trainer.logger = logger
 
+        # connect trainer to logger
+        logger.connect(self.trainer)
+
     def cache_training_step_metrics(self, opt_closure_result):
         """
         This function is responsible to update