Fix model tests: unique directory for checkpoints

sktime · Aug 15, 2020 · e909fc4 · e909fc4
1 parent 53ab0b2
commit e909fc4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/tests/test_models/test_nbeats.py b/tests/test_models/test_nbeats.py
@@ -1,7 +1,7 @@
 import shutil
 import pytorch_lightning as pl
 from pytorch_lightning.loggers import TensorBoardLogger
-from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 from pytorch_forecasting.metrics import QuantileLoss
 from pytorch_forecasting.models import NBeats
 
@@ -12,7 +12,9 @@ def test_integration(dataloaders_fixed_window_without_coveratiates, tmp_path):
     early_stop_callback = EarlyStopping(monitor="val_loss", min_delta=1e-4, patience=1, verbose=False, mode="min")
 
     logger = TensorBoardLogger(tmp_path)
+    checkpoint = ModelCheckpoint(filepath=tmp_path)
     trainer = pl.Trainer(
+        checkpoint_callback=checkpoint,
         max_epochs=3,
         gpus=0,
         weights_summary="top",
@@ -28,6 +30,12 @@ def test_integration(dataloaders_fixed_window_without_coveratiates, tmp_path):
         trainer.fit(
             net, train_dataloader=train_dataloader, val_dataloaders=val_dataloader,
         )
+        # check loading
+        fname = f"{trainer.checkpoint_callback.dirpath}/epoch=0.ckpt"
+        net = NBeats.load_from_checkpoint(fname)
+
+        # check prediction
+        net.predict(val_dataloader, fast_dev_run=True, return_index=True, return_decoder_lengths=True)
     finally:
         shutil.rmtree(tmp_path, ignore_errors=True)
 

diff --git a/tests/test_models/test_temporal_fusion_transformer.py b/tests/test_models/test_temporal_fusion_transformer.py
@@ -1,7 +1,7 @@
 import shutil
 import pytorch_lightning as pl
 from pytorch_lightning.loggers import TensorBoardLogger
-from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 from pytorch_forecasting.metrics import QuantileLoss
 from pytorch_forecasting.models import TemporalFusionTransformer
 
@@ -17,7 +17,9 @@ def test_integration(dataloaders_with_coveratiates, tmp_path):
 
     # check training
     logger = TensorBoardLogger(tmp_path)
+    checkpoint = ModelCheckpoint(filepath=tmp_path)
     trainer = pl.Trainer(
+        checkpoint_callback=checkpoint,
         max_epochs=3,
         gpus=0,
         weights_summary="top",