fixing tested values

Author: Borda

tests/helpers/pipelines.py

@@ -13,19 +13,23 @@
 # limitations under the License.
 import torch
 
-from pytorch_lightning import Trainer
+from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+from pytorch_lightning.metrics.functional import accuracy
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.utilities import DistributedType
+from tests.base import EvalModelTemplate
 from tests.helpers import BoringModel
 from tests.helpers.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):
+def run_model_test_without_loggers(
+    trainer_options: dict, model: LightningModule, data: LightningDataModule = None, min_acc: float = 0.50
+):
     reset_seed()
 
     # fit model
     trainer = Trainer(**trainer_options)
-    trainer.fit(model)
+    trainer.fit(model, datamodule=data)
 
     # correct result and ok accuracy
     assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
@@ -35,12 +39,13 @@ def run_model_test_without_loggers(trainer_options, model, min_acc: float = 0.50
     )
 
     # test new model accuracy
-    test_loaders = model.test_dataloader()
+    test_loaders = model.test_dataloader() if not data else data.test_dataloader()
     if not isinstance(test_loaders, list):
         test_loaders = [test_loaders]
 
-    for dataloader in test_loaders:
-        run_prediction(pretrained_model, dataloader, min_acc=min_acc)
+    if not isinstance(model, BoringModel):
+        for dataloader in test_loaders:
+            run_prediction_eval_model_template(model, dataloader, min_acc=min_acc)
 
     if trainer._distrib_type in (DistributedType.DDP, DistributedType.DDP_SPAWN):
         # on hpc this would work fine... but need to hack it for the purpose of the test
@@ -77,8 +82,9 @@ def run_model_test(
     if not isinstance(test_loaders, list):
         test_loaders = [test_loaders]
 
-    for dataloader in test_loaders:
-        run_prediction(pretrained_model, dataloader, min_acc=min_acc)
+    if not isinstance(model, BoringModel):
+        for dataloader in test_loaders:
+            run_prediction_eval_model_template(model, dataloader, min_acc=min_acc)
 
     if with_hpc:
         if trainer._distrib_type in (DistributedType.DDP, DistributedType.DDP_SPAWN, DistributedType.DDP2):
@@ -95,14 +101,7 @@ def run_model_test(
         trainer.checkpoint_connector.hpc_load(checkpoint_path, on_gpu=on_gpu)
 
 
-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):
+def run_prediction_eval_model_template(trained_model, dataloader, dp=False, min_acc=0.50):
     # run prediction on 1 batch
     batch = next(iter(dataloader))
     x, y = batch
@@ -117,24 +116,6 @@ def _eval_model_template_run_prediction(trained_model, dataloader, dp=False, min
     else:
         with torch.no_grad():
             y_hat = trained_model(x)
-        y_hat = y_hat.cpu()
-
-        # acc
-        labels_hat = torch.argmax(y_hat, dim=1)
-
-        y = y.cpu()
-        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-        acc = torch.tensor(acc)
-        acc = acc.item()
+        acc = accuracy(y_hat.cpu(), y.cpu()).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}"

tests/models/data/horovod/train_default_model.py

@@ -37,7 +37,7 @@
     print('You requested to import Horovod which is missing or not supported for your OS.')
 
 from tests.base import EvalModelTemplate  # noqa: E402
-from tests.helpers.pipelines import run_prediction  # noqa: E402
+from tests.helpers.pipelines import run_prediction_eval_model_template  # noqa: E402
 from tests.helpers.utils import reset_seed, set_random_master_port  # noqa: E402
 
 parser = argparse.ArgumentParser()
@@ -74,7 +74,7 @@ def run_test_from_config(trainer_options):
         test_loaders = [test_loaders]
 
     for dataloader in test_loaders:
-        run_prediction(pretrained_model, dataloader)
+        run_prediction_eval_model_template(pretrained_model, dataloader)
 
     # test HPC saving
     trainer.checkpoint_connector.hpc_save(ckpt_path, trainer.logger)

tests/models/test_cpu.py

@@ -25,6 +25,8 @@
 from pytorch_lightning.trainer.states import TrainerState
 from tests.base import EvalModelTemplate
 from tests.helpers import BoringModel
+from tests.helpers.datamodules import ClassifDataModule
+from tests.helpers.simple_models import ClassificationModel
 
 
 def test_cpu_slurm_save_load(tmpdir):
@@ -149,7 +151,7 @@ def test_multi_cpu_model_ddp(tmpdir):
 def test_lbfgs_cpu_model(tmpdir):
     """Test each of the trainer options. Testing LBFGS optimizer"""
 
-    class ModelSpecifiedOptimizer(BoringModel):
+    class ModelSpecifiedOptimizer(ClassificationModel):
 
         def __init__(self, optimizer_name, learning_rate):
             super().__init__()
@@ -166,8 +168,9 @@ def __init__(self, optimizer_name, learning_rate):
         limit_val_batches=0.2,
     )
 
+    dm = ClassifDataModule()
     model = ModelSpecifiedOptimizer(optimizer_name="LBFGS", learning_rate=0.004)
-    tpipes.run_model_test_without_loggers(trainer_options, model, min_acc=0.25)
+    tpipes.run_model_test_without_loggers(trainer_options, model, dm, min_acc=0.25)
 
 
 def test_default_logger_callbacks_cpu_model(tmpdir):

tests/models/test_gpu.py

@@ -25,6 +25,7 @@
 from pytorch_lightning.utilities import device_parser
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
+from tests.helpers.simple_models import ClassificationModel
 
 PRETEND_N_OF_GPUS = 16
 
@@ -42,8 +43,8 @@ def test_multi_gpu_none_backend(tmpdir):
         gpus=2,
     )
 
-    model = BoringModel()
-    tpipes.run_model_test(trainer_options, model, min_acc=0.20)
+    model = ClassificationModel()
+    tpipes.run_model_test(trainer_options, model)
 
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")

tests/models/test_horovod.py

@@ -33,6 +33,8 @@
 from tests.base import EvalModelTemplate
 from tests.helpers.advanced_models import BasicGAN
 from tests.helpers.boring_model import BoringModel
+from tests.helpers.datamodules import ClassifDataModule
+from tests.helpers.simple_models import ClassificationModel
 
 if _HOROVOD_AVAILABLE:
     import horovod
@@ -173,7 +175,7 @@ def test_horovod_amp(tmpdir):
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
 def test_horovod_transfer_batch_to_gpu(tmpdir):
 
-    class TestTrainingStepModel(EvalModelTemplate):
+    class TestTrainingStepModel(ClassificationModel):
 
         def training_step(self, batch, *args, **kwargs):
             x, y = batch
@@ -200,7 +202,9 @@ def validation_step(self, batch, *args, **kwargs):
         deterministic=True,
         accelerator='horovod',
     )
-    tpipes.run_model_test_without_loggers(trainer_options, model)
+
+    dm = ClassifDataModule()
+    tpipes.run_model_test_without_loggers(trainer_options, model, dm, min_acc=0.25)
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="Horovod is not supported on Windows")

tests/models/test_restore.py

@@ -200,7 +200,7 @@ def test_running_test_pretrained_model_distrib_dp(tmpdir):
         dataloaders = [dataloaders]
 
     for dataloader in dataloaders:
-        tpipes.run_prediction(pretrained_model, dataloader)
+        tpipes.run_prediction_eval_model_template(pretrained_model, dataloader)
 
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
@@ -251,7 +251,7 @@ def test_running_test_pretrained_model_distrib_ddp_spawn(tmpdir):
         dataloaders = [dataloaders]
 
     for dataloader in dataloaders:
-        tpipes.run_prediction(pretrained_model, dataloader)
+        tpipes.run_prediction_eval_model_template(pretrained_model, dataloader)
 
 
 def test_running_test_pretrained_model_cpu(tmpdir):
@@ -394,7 +394,7 @@ def assert_good_acc():
         dp_model.module.module.running_stage = RunningStage.EVALUATING
 
         dataloader = trainer.train_dataloader
-        tpipes.run_prediction(dp_model, dataloader, dp=True)
+        tpipes.run_prediction_eval_model_template(dp_model, dataloader, dp=True)
 
     # new model
     model = EvalModelTemplate(**hparams)

tests/trainer/test_dataloaders.py

@@ -131,7 +131,7 @@ def test_multiple_val_dataloader(tmpdir):
 
     # make sure predictions are good for each val set
     for dataloader in trainer.val_dataloaders:
-        tpipes.run_prediction(trainer.model, dataloader)
+        tpipes.run_prediction_eval_model_template(trainer.model, dataloader)
 
 
 @pytest.mark.parametrize('ckpt_path', [None, 'best', 'specific'])
@@ -168,7 +168,7 @@ def test_step(self, batch, batch_idx, *args, **kwargs):
 
     # make sure predictions are good for each test set
     for dataloader in trainer.test_dataloaders:
-        tpipes.run_prediction(trainer.model, dataloader)
+        tpipes.run_prediction_eval_model_template(trainer.model, dataloader)
 
     # run the test method
     trainer.test(ckpt_path=ckpt_path)