rm EvalModelTemplate

Author: rohitgr7

tests/trainer/flags/test_limit_batches.py

@@ -0,0 +1,75 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.trainer.states import RunningStage
+from tests.helpers.boring_model import BoringModel
+
+
+def test_num_dataloader_batches(tmpdir):
+    """Tests that the correct number of batches are allocated."""
+    # when we have fewer batches in the dataloader we should use those instead of the limit
+    model = BoringModel()
+    trainer = Trainer(limit_val_batches=100, limit_train_batches=100, max_epochs=1, default_root_dir=tmpdir)
+    trainer.fit(model)
+
+    assert len(model.train_dataloader()) == 64
+    assert len(model.val_dataloader()) == 64
+    assert isinstance(trainer.num_val_batches, list)
+    assert trainer.num_val_batches[0] == 64
+    assert trainer.num_training_batches == 64
+
+    # when we have more batches in the dataloader we should limit them
+    model = BoringModel()
+    trainer = Trainer(limit_val_batches=7, limit_train_batches=7, max_epochs=1, default_root_dir=tmpdir)
+    trainer.fit(model)
+
+    assert len(model.train_dataloader()) == 64
+    assert len(model.val_dataloader()) == 64
+    assert isinstance(trainer.num_val_batches, list)
+    assert trainer.num_val_batches[0] == 7
+    assert trainer.num_training_batches == 7
+
+
+@pytest.mark.parametrize(
+    ["stage", "mode"],
+    [
+        (RunningStage.VALIDATING, "val"),
+        (RunningStage.TESTING, "test"),
+        (RunningStage.PREDICTING, "predict"),
+    ],
+)
+def test_eval_limit_batches(stage, mode):
+    limit_eval_batches = f"limit_{mode}_batches"
+    dl_hook = f"{mode}_dataloader"
+    model = BoringModel()
+    eval_loader = getattr(model, dl_hook)()
+
+    limit_batches = 0.1
+    trainer = Trainer(**{limit_eval_batches: limit_batches})
+    model.trainer = trainer
+    trainer._data_connector.attach_dataloaders(model)
+    loader_num_batches, dataloaders = trainer._reset_eval_dataloader(stage, model=model)
+    assert loader_num_batches[0] == int(limit_batches * len(eval_loader))
+    assert len(dataloaders[0]) == len(eval_loader)
+
+    limit_batches = 10
+    trainer = Trainer(**{limit_eval_batches: limit_batches})
+    model.trainer = trainer
+    trainer._data_connector.attach_dataloaders(model)
+    loader_num_batches, dataloaders = trainer._reset_eval_dataloader(stage, model=model)
+    assert loader_num_batches[0] == limit_batches
+    assert len(dataloaders[0]) == len(eval_loader)

tests/trainer/flags/test_overfit_batches.py

@@ -15,9 +15,9 @@
 import torch
 from torch.utils.data import DataLoader, RandomSampler, Sampler, SequentialSampler
 
+from legacy.simple_classif_training import ClassifDataModule, ClassificationModel
 from pytorch_lightning import Trainer
 from pytorch_lightning.trainer.states import RunningStage
-from tests.base.model_template import EvalModelTemplate
 from tests.helpers.boring_model import BoringModel, RandomDataset
 
 
@@ -70,32 +70,32 @@ def test_overfit_batch_limits(tmpdir):
     # ------------------------------------------------------
     # Make sure shuffle is correct across loaders initially
     # ------------------------------------------------------
-    model = EvalModelTemplate()
-    model.train_dataloader()
+    model = ClassificationModel()
+    dm = ClassifDataModule()
 
     # original train loader which should be replaced in all methods
-    train_loader = model.train_dataloader()
+    train_loader = dm.train_dataloader()
 
     # make sure the val and tests are not shuffled
     assert isinstance(train_loader.sampler, RandomSampler)
-    assert isinstance(model.val_dataloader().sampler, SequentialSampler)
-    assert isinstance(model.test_dataloader().sampler, SequentialSampler)
+    assert isinstance(dm.val_dataloader().sampler, SequentialSampler)
+    assert isinstance(dm.test_dataloader().sampler, SequentialSampler)
 
     # ------------------------------------------------------
     # get the training loader and batch
     # ------------------------------------------------------
     # Create a reference train dataloader without shuffling.
-    train_loader = DataLoader(model.train_dataloader().dataset, shuffle=False)
+    train_loader = DataLoader(dm.train_dataloader().dataset, shuffle=False)
     (xa, ya) = next(iter(train_loader))
-    train_loader = DataLoader(model.train_dataloader().dataset, shuffle=True)
+    train_loader = DataLoader(dm.train_dataloader().dataset, shuffle=True)
     full_train_samples = len(train_loader)
     num_train_samples = int(0.11 * full_train_samples)
 
     # ------------------------------------------------------
     # set VAL and Test loaders
     # ------------------------------------------------------
-    val_loader = DataLoader(model.val_dataloader().dataset, shuffle=False)
-    test_loader = DataLoader(model.test_dataloader().dataset, shuffle=False)
+    val_loader = DataLoader(dm.val_dataloader().dataset, shuffle=False)
+    test_loader = DataLoader(dm.test_dataloader().dataset, shuffle=False)
 
     # set the model loaders
     model.train_dataloader = lambda: train_loader
@@ -165,27 +165,3 @@ def test_overfit_batch_limits(tmpdir):
             assert loader_num_batches[0] == 0
         else:
             assert loader_num_batches[0] == len(test_loader)
-
-        # ------------------------------------------------------
-        # test limit_xxx_batches as percent AND int
-        # ------------------------------------------------------
-        if split == RunningStage.VALIDATING:
-            trainer = Trainer(limit_val_batches=0.1)
-            trainer._data_connector.attach_dataloaders(model)
-            loader_num_batches, dataloaders = trainer._reset_eval_dataloader(split, model=model)
-            assert loader_num_batches[0] == int(0.1 * len(val_loader))
-
-            trainer = Trainer(limit_val_batches=10)
-            trainer._data_connector.attach_dataloaders(model)
-            loader_num_batches, dataloaders = trainer._reset_eval_dataloader(split, model=model)
-            assert loader_num_batches[0] == 10
-        else:
-            trainer = Trainer(limit_test_batches=0.1)
-            trainer._data_connector.attach_dataloaders(model)
-            loader_num_batches, dataloaders = trainer._reset_eval_dataloader(split, model=model)
-            assert loader_num_batches[0] == int(0.1 * len(test_loader))
-
-            trainer = Trainer(limit_test_batches=10)
-            trainer._data_connector.attach_dataloaders(model)
-            loader_num_batches, dataloaders = trainer._reset_eval_dataloader(split, model=model)
-            assert loader_num_batches[0] == 10