Update tests to reflect changed reload_dataloaders_every_n_epochs behavior

Author: adamviola

tests/models/test_hooks.py

@@ -872,7 +872,6 @@ def call(hook, fn, *args, **kwargs):
         dict(name="setup", kwargs=dict(stage="fit")),
         dict(name="val_dataloader"),
         dict(name="train_dataloader"),
-        dict(name="val_dataloader"),
         dict(name="on_save_checkpoint", args=(ANY,)),
         dict(name="teardown", kwargs=dict(stage="fit")),
     ]

tests/trainer/test_dataloaders.py

@@ -1118,52 +1118,6 @@ def test_dataloaders_load_only_once(tmpdir):
 
     assert tracker.mock_calls == [call.val_dataloader(), call.train_dataloader()]
 
-
-def test_dataloaders_load_only_once_val_interval(tmpdir):
-    model = BoringModel()
-
-    # logger file to get meta
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        limit_train_batches=10,
-        limit_val_batches=10,
-        val_check_interval=0.3,
-        reload_dataloaders_every_n_epochs=1,
-        max_epochs=3,
-    )
-
-    tracker = Mock()
-    model.train_dataloader = Mock(wraps=model.train_dataloader)
-    model.val_dataloader = Mock(wraps=model.val_dataloader)
-    model.test_dataloader = Mock(wraps=model.test_dataloader)
-
-    tracker.attach_mock(model.train_dataloader, "train_dataloader")
-    tracker.attach_mock(model.val_dataloader, "val_dataloader")
-    tracker.attach_mock(model.test_dataloader, "test_dataloader")
-
-    trainer.fit(model)
-    trainer.test(model)
-
-    # verify the sequence
-    expected_sequence = [
-        call.val_dataloader(),
-        call.train_dataloader(),
-        call.val_dataloader(),
-        call.val_dataloader(),
-        call.val_dataloader(),
-        call.train_dataloader(),
-        call.val_dataloader(),
-        call.val_dataloader(),
-        call.val_dataloader(),
-        call.train_dataloader(),
-        call.val_dataloader(),
-        call.val_dataloader(),
-        call.val_dataloader(),
-        call.test_dataloader(),
-    ]
-    assert tracker.mock_calls == expected_sequence
-
-
 def test_dataloaders_load_only_once_no_sanity_check(tmpdir):
     model = BoringModel()
 
@@ -1190,14 +1144,25 @@ def test_dataloaders_load_only_once_no_sanity_check(tmpdir):
 
 @pytest.mark.parametrize("n", [1, 2])
 def test_dataloaders_load_every_n_epochs(tmpdir, n):
-    model = BoringModel()
+    train_reload_epochs, val_reload_epochs = [], []
+
+    class TestModel(BoringModel):
+        def train_dataloader(self):
+            train_reload_epochs.append(self.current_epoch)
+            return super().train_dataloader()
+
+        def val_dataloader(self):
+            val_reload_epochs.append(self.current_epoch)
+            return super().val_dataloader()
+    
+    model = TestModel()
 
     trainer = Trainer(
         default_root_dir=tmpdir,
         limit_train_batches=0.3,
         limit_val_batches=0.3,
         reload_dataloaders_every_n_epochs=n,
-        max_epochs=3,
+        max_epochs=5,
     )
 
     tracker = Mock()
@@ -1212,15 +1177,24 @@ def test_dataloaders_load_every_n_epochs(tmpdir, n):
     trainer.fit(model)
     trainer.test(model)
 
-    # verify the sequence
-    expected_sequence = [call.val_dataloader()]
+    # Verify the sequence
+    expected_sequence = [call.val_dataloader(), call.train_dataloader()] # Sanity check first
     if n == 1:
-        expected_sequence += [call.train_dataloader(), call.val_dataloader()] * 3
+        expected_sequence += [call.train_dataloader(), call.val_dataloader()] * 4
     elif n == 2:
         expected_sequence += [call.train_dataloader(), call.val_dataloader()] * 2
     expected_sequence += [call.test_dataloader()]
+
     assert tracker.mock_calls == expected_sequence
 
+    # Verify epoch of reloads
+    if n == 1:
+        assert train_reload_epochs == [0, 1, 2, 3, 4]
+        assert val_reload_epochs == [0, 1, 2, 3, 4]
+    elif n == 2:
+        assert train_reload_epochs == [0, 2, 4]
+        assert val_reload_epochs == [0, 2, 4]
+
 
 @pytest.mark.parametrize("n", [3, 5])
 def test_dataloaders_load_every_n_epochs_infrequent_val(tmpdir, n):
@@ -1263,6 +1237,7 @@ def val_dataloader(self):
 
     model.test_dataloader.assert_called_once()
 
+
 def test_dataloaders_load_every_n_epochs_frequent_val(tmpdir):
     train_reload_epochs, val_reload_epochs, val_check_epochs = [], [], []
 
@@ -1303,6 +1278,7 @@ def validation_epoch_end(self, outputs):
     # Verify validation happens 3 times per epoch + 1 for sanity check
     assert val_check_epochs == [0, 0, 0, 0, 1, 1, 1, 2, 2, 2]
 
+
 @pytest.mark.parametrize("n", ["test", -1])
 def test_dataloaders_load_every_n_epochs_exception(tmpdir, n):
 
@@ -1347,15 +1323,6 @@ def validation_step(self, batch, batch_idx):
     expected_calls = [
         call.train_dataloader(),
         call.val_dataloader(),
-        # This has subsequent calls to val_dataloader
-        # because the training loop runs the evaluation loop,
-        # which reloads the val dataloader again.
-        # We cannot yet rely on trainer.current_epoch=0 to skip reloading
-        # the val dataloader on the first epoch because this only tracks the training epoch
-        # meaning multiple passes through the validation data within a single training epoch
-        # would not have the dataloader reloaded.
-        # This breaks the assumption behind reload_dataloaders_every_n_epochs=1
-        call.val_dataloader(),
         call.train_dataloader(),
         call.val_dataloader(),
         call.train_dataloader(),