save learning plot

pipelines/trainer.py

@@ -11,6 +11,7 @@
 from core.constants import DEVICE
 from pipelines.base import BaseRunner
 from pipelines.utils.early_stopping import EarlyStopping
+from pipelines.utils.visualize_utils import save_learning_curve_plot
 
 
 class TrainingParams(TypedDict):
@@ -83,7 +84,7 @@ def run(self) -> None:
                 print(f"Early stopped at epoch {epoch}")
                 break
 
-        self.__save_metrics()
+        self._save_artifacts()
 
     @property
     def training_metrics(self) -> TrainingMetrics:
@@ -139,7 +140,12 @@ def __log_metric(
     def __latest_training_metric(self) -> Dict[str, float]:
         return {k: cast(List[float], v)[-1] for k, v in self._training_metrics.items()}
 
-    def __save_metrics(self) -> None:
+    def _save_artifacts(self) -> None:
         pd.DataFrame(self._training_metrics).to_csv(
             os.path.join(self.artifact_dir, self.metrics_filename)
         )
+        save_learning_curve_plot(
+            os.path.join(self.artifact_dir, "learning_curve.png"),
+            self._training_metrics["train_loss"],
+            self._training_metrics["validation_loss"],
+        )

pipelines/utils/visualize_utils.py

@@ -11,8 +11,7 @@
 
 
 def save_learning_curve_plot(
-    save_dir_path: str,
-    model_name: str,
+    save_img_path: str,
     training_losses: List,
     validation_losses: List,
 ) -> None:
@@ -47,8 +46,7 @@ def save_learning_curve_plot(
 
     ax.legend(loc="upper center")
     plt.tight_layout()
-    save_path = os.path.join(save_dir_path, f"{model_name}_training_results.png")
-    plt.savefig(save_path)
+    plt.savefig(save_img_path)
     plt.close()
 
 

tests/pipelines/test_experimenter.py

@@ -46,6 +46,7 @@ def test_run(mocked_save_seq2seq_model):
         # testing trainer artifacts
         assert os.path.exists(os.path.join(tempdirpath, "train", "model.pt"))
         assert os.path.exists(os.path.join(tempdirpath, "train", "metrics.csv"))
+        assert os.path.exists(os.path.join(tempdirpath, "train", "learning_curve.png"))
         # testing evaluator artifacts
         for i in range(dataset_length):
             assert os.path.exists(

tests/pipelines/test_trainer.py

@@ -43,6 +43,7 @@ def test_run(mocked_save_seq2seq_model):
 
         assert os.path.exists(os.path.join(tempdirpath, "checkpoint.pt"))
         assert os.path.exists(os.path.join(tempdirpath, "example.csv"))
+        assert os.path.exists(os.path.join(tempdirpath, "learning_curve.png"))
         for metrics in trainer.training_metrics.values():
             assert len(metrics) == epochs
 
@@ -75,5 +76,6 @@ def test_run_early_stopping(mocked_save_seq2seq_model):
 
         assert os.path.exists(os.path.join(tempdirpath, "checkpoint.pt"))
         assert os.path.exists(os.path.join(tempdirpath, "example.csv"))
+        assert os.path.exists(os.path.join(tempdirpath, "learning_curve.png"))
         for metrics in trainer.training_metrics.values():
             assert len(metrics) == epochs - patience