Lightning-AI
diff --git a/‎tests/base/develop_pipelines.py‎
Lines changed: 32 additions & 14 deletions b/‎tests/base/develop_pipelines.py‎
Lines changed: 32 additions & 14 deletions
diff --git a/‎tests/models/data/horovod/train_default_model.py‎
Lines changed: 1 addition & 1 deletion b/‎tests/models/data/horovod/train_default_model.py‎
Lines changed: 1 addition & 1 deletion
@@ -14,8 +14,8 @@
 import torch
 
 from pytorch_lightning import Trainer
-from tests.base.develop_utils import load_model_from_checkpoint, get_default_logger, \
-    reset_seed
+from tests.base import BoringModel
+from tests.base.develop_utils import get_default_logger, load_model_from_checkpoint, reset_seed
 
 
 def run_model_test_without_loggers(trainer_options, model, min_acc: float = 0.50):
@@ -31,6 +31,7 @@ def run_model_test_without_loggers(trainer_options, model, min_acc: float = 0.50
     pretrained_model = load_model_from_checkpoint(
         trainer.logger,
         trainer.checkpoint_callback.best_model_path,
+        type(model)
     )
 
     # test new model accuracy
@@ -39,15 +40,16 @@ def run_model_test_without_loggers(trainer_options, model, min_acc: float = 0.50
         test_loaders = [test_loaders]
 
     for dataloader in test_loaders:
-        run_prediction(dataloader, pretrained_model, min_acc=min_acc)
+        run_prediction(pretrained_model, dataloader, min_acc=min_acc)
 
     if trainer.use_ddp:
         # on hpc this would work fine... but need to hack it for the purpose of the test
         trainer.model = pretrained_model
         trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
 
 
-def run_model_test(trainer_options, model, on_gpu: bool = True, version=None, with_hpc: bool = True):
+def run_model_test(trainer_options, model, on_gpu: bool = True, version=None,
+                   with_hpc: bool = True, min_acc: float = 0.25):
 
     reset_seed()
     save_dir = trainer_options['default_root_dir']
@@ -56,35 +58,34 @@ def run_model_test(trainer_options, model, on_gpu: bool = True, version=None, wi
     logger = get_default_logger(save_dir, version=version)
     trainer_options.update(logger=logger)
 
-    if 'checkpoint_callback' not in trainer_options:
-        trainer_options.update(checkpoint_callback=True)
-
     trainer = Trainer(**trainer_options)
     initial_values = torch.tensor([torch.sum(torch.abs(x)) for x in model.parameters()])
     result = trainer.fit(model)
     post_train_values = torch.tensor([torch.sum(torch.abs(x)) for x in model.parameters()])
 
     assert result == 1, 'trainer failed'
     # Check that the model is actually changed post-training
-    assert torch.norm(initial_values - post_train_values) > 0.1
+    change_ratio = torch.norm(initial_values - post_train_values)
+    assert change_ratio > 0.1, f"the model is changed of {change_ratio}"
 
     # test model loading
-    pretrained_model = load_model_from_checkpoint(logger, trainer.checkpoint_callback.best_model_path)
+    pretrained_model = load_model_from_checkpoint(logger, trainer.checkpoint_callback.best_model_path, type(model))
 
     # test new model accuracy
     test_loaders = model.test_dataloader()
     if not isinstance(test_loaders, list):
         test_loaders = [test_loaders]
 
     for dataloader in test_loaders:
-        run_prediction(dataloader, pretrained_model)
+        run_prediction(pretrained_model, dataloader, min_acc=min_acc)
 
     if with_hpc:
         if trainer.use_ddp or trainer.use_ddp2:
             # on hpc this would work fine... but need to hack it for the purpose of the test
             trainer.model = pretrained_model
-            trainer.optimizers, trainer.lr_schedulers, trainer.optimizer_frequencies = \
-                trainer.init_optimizers(pretrained_model)
+            trainer.optimizers, trainer.lr_schedulers, trainer.optimizer_frequencies = trainer.init_optimizers(
+                pretrained_model
+            )
 
         # test HPC saving
         trainer.checkpoint_connector.hpc_save(save_dir, logger)
@@ -93,7 +94,14 @@ def run_model_test(trainer_options, model, on_gpu: bool = True, version=None, wi
         trainer.checkpoint_connector.hpc_load(checkpoint_path, on_gpu=on_gpu)
 
 
-def run_prediction(dataloader, trained_model, dp=False, min_acc=0.50):
+def run_prediction(trained_model, dataloader, dp=False, min_acc=0.25):
+    if isinstance(trained_model, BoringModel):
+        return _boring_model_run_prediction(trained_model, dataloader, dp, min_acc)
+    else:
+        return _eval_model_template_run_prediction(trained_model, dataloader, dp, min_acc)
+
+
+def _eval_model_template_run_prediction(trained_model, dataloader, dp=False, min_acc=0.50):
     # run prediction on 1 batch
     batch = next(iter(dataloader))
     x, y = batch
@@ -102,7 +110,7 @@ def run_prediction(dataloader, trained_model, dp=False, min_acc=0.50):
     if dp:
         with torch.no_grad():
             output = trained_model(batch, 0)
-        acc = output['val_acc']
+            acc = output['val_acc']
         acc = torch.mean(acc).item()
 
     else:
@@ -119,3 +127,13 @@ def run_prediction(dataloader, trained_model, dp=False, min_acc=0.50):
         acc = acc.item()
 
     assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"
+
+
+def _boring_model_run_prediction(trained_model, dataloader, dp=False, min_acc=0.25):
+    # run prediction on 1 batch
+    batch = next(iter(dataloader))
+    with torch.no_grad():
+        output = trained_model(batch)
+    acc = trained_model.loss(batch, output)
+
+    assert acc >= min_acc, f"This model is expected to get, {min_acc} in test set but got {acc}"
@@ -72,7 +72,7 @@ def run_test_from_config(trainer_options):
         test_loaders = [test_loaders]
 
     for dataloader in test_loaders:
-        run_prediction(dataloader, pretrained_model)
+        run_prediction(pretrained_model, dataloader)
 
     # test HPC saving
     trainer.checkpoint_connector.hpc_save(ckpt_path, trainer.logger)