From 4ac9925dad71edda26db486795341f0b0ba4ed38 Mon Sep 17 00:00:00 2001
From: fdelrio89 <fdelrio89@gmail.com>
Date: Fri, 21 Feb 2020 22:47:48 -0300
Subject: [PATCH] Fix comet logger to log after train (#892)

* Fix comet logger to log after train

* Add clarifying comment to COmetLogger code

Explains the need to use CometExistingExperiment in the CometLogger class after
CometLogger.finalize.
---
 pytorch_lightning/loggers/comet.py | 39 ++++++++++++++++++++++++------
 1 file changed, 32 insertions(+), 7 deletions(-)

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):