diff --git a/pytorch_lightning/loggers/comet.py b/pytorch_lightning/loggers/comet.py index 06b2714873e94..fd8d456f34f41 100644 --- a/pytorch_lightning/loggers/comet.py +++ b/pytorch_lightning/loggers/comet.py @@ -10,6 +10,7 @@ try: from comet_ml import Experiment as CometExperiment + from comet_ml import ExistingExperiment as CometExistingExperiment from comet_ml import OfflineExperiment as CometOfflineExperiment try: from comet_ml.api import API @@ -33,7 +34,8 @@ class CometLogger(LightningLoggerBase): """ def __init__(self, api_key=None, save_dir=None, workspace=None, - rest_api_key=None, project_name=None, experiment_name=None, **kwargs): + rest_api_key=None, project_name=None, experiment_name=None, + experiment_key=None, **kwargs): r""" Requires either an API Key (online mode) or a local directory path (offline mode) @@ -100,6 +102,7 @@ def __init__(self, api_key=None, save_dir=None, workspace=None, self.workspace = workspace self.project_name = project_name + self.experiment_key = experiment_key self._kwargs = kwargs if rest_api_key is not None: @@ -131,12 +134,22 @@ def experiment(self): return self._experiment if self.mode == "online": - self._experiment = CometExperiment( - api_key=self.api_key, - workspace=self.workspace, - project_name=self.project_name, - **self._kwargs - ) + if self.experiment_key is None: + self._experiment = CometExperiment( + api_key=self.api_key, + workspace=self.workspace, + project_name=self.project_name, + **self._kwargs + ) + self.experiment_key = self._experiment.get_key() + else: + self._experiment = CometExistingExperiment( + api_key=self.api_key, + workspace=self.workspace, + project_name=self.project_name, + previous_experiment=self.experiment_key, + **self._kwargs + ) else: self._experiment = CometOfflineExperiment( offline_directory=self.save_dir, @@ -160,9 +173,21 @@ def log_metrics(self, metrics, step=None): self.experiment.log_metrics(metrics, step=step) + def reset_experiment(self): + self._experiment = None + @rank_zero_only def finalize(self, status): + r""" + When calling self.experiment.end(), that experiment won't log any more data to Comet. That's why, if you need + to log any more data you need to create an ExistingCometExperiment. For example, to log data when testing your + model after training, because when training is finalized CometLogger.finalize is called. + + This happens automatically in the CometLogger.experiment property, when self._experiment is set to None + i.e. self.reset_experiment(). + """ self.experiment.end() + self.reset_experiment() @property def name(self):