Speeding up and reducing the memory footprint of the trainer tests

tests/trainers/test_byol.py

@@ -14,6 +14,8 @@
 from torchgeo.trainers import BYOLTask
 from torchgeo.trainers.byol import BYOL, SimCLRAugmentation
 
+from .test_utils import ClassificationTestModel
+
 
 class TestBYOL:
     def test_custom_augment_fn(self) -> None:
@@ -53,6 +55,8 @@ def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
         model_kwargs = conf_dict["module"]
         model = BYOLTask(**model_kwargs)
 
+        model.encoder = ClassificationTestModel(**model_kwargs)
+
         # Instantiate trainer
         trainer = Trainer(fast_dev_run=True, log_every_n_steps=1)
         trainer.fit(model=model, datamodule=datamodule)

tests/trainers/test_classification.py

@@ -2,11 +2,14 @@
 # Licensed under the MIT License.
 
 import os
-from typing import Any, Dict, Type, cast
+from typing import Any, Dict, Generator, Type, cast
 
 import pytest
+import timm
+from _pytest.monkeypatch import MonkeyPatch
 from omegaconf import OmegaConf
 from pytorch_lightning import LightningDataModule, Trainer
+from torch.nn.modules import Module
 
 from torchgeo.datamodules import (
     BigEarthNetDataModule,
@@ -17,6 +20,12 @@
 )
 from torchgeo.trainers import ClassificationTask, MultiLabelClassificationTask
 
+from .test_utils import ClassificationTestModel
+
+
+def create_model(*args: Any, **kwargs: Any) -> Module:
+    return ClassificationTestModel(**kwargs)
+
 
 class TestClassificationTask:
     @pytest.mark.parametrize(
@@ -29,7 +38,12 @@ class TestClassificationTask:
             ("ucmerced", UCMercedDataModule),
         ],
     )
-    def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
+    def test_trainer(
+        self,
+        monkeypatch: Generator[MonkeyPatch, None, None],
+        name: str,
+        classname: Type[LightningDataModule],
+    ) -> None:
         if name.startswith("so2sat"):
             pytest.importorskip("h5py")
 
@@ -42,6 +56,9 @@ def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
         datamodule = classname(**datamodule_kwargs)
 
         # Instantiate model
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            timm, "create_model", create_model
+        )
         model_kwargs = conf_dict["module"]
         model = ClassificationTask(**model_kwargs)
 
@@ -119,7 +136,12 @@ class TestMultiLabelClassificationTask:
             ("bigearthnet_s2", BigEarthNetDataModule),
         ],
     )
-    def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
+    def test_trainer(
+        self,
+        monkeypatch: Generator[MonkeyPatch, None, None],
+        name: str,
+        classname: Type[LightningDataModule],
+    ) -> None:
         conf = OmegaConf.load(os.path.join("conf", "task_defaults", name + ".yaml"))
         conf_dict = OmegaConf.to_object(conf.experiment)
         conf_dict = cast(Dict[Any, Dict[Any, Any]], conf_dict)
@@ -129,6 +151,9 @@ def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
         datamodule = classname(**datamodule_kwargs)
 
         # Instantiate model
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            timm, "create_model", create_model
+        )
         model_kwargs = conf_dict["module"]
         model = MultiLabelClassificationTask(**model_kwargs)
 

tests/trainers/test_regression.py

@@ -11,6 +11,8 @@
 from torchgeo.datamodules import COWCCountingDataModule, CycloneDataModule
 from torchgeo.trainers import RegressionTask
 
+from .test_utils import RegressionTestModel
+
 
 class TestRegressionTask:
     @pytest.mark.parametrize(
@@ -30,6 +32,8 @@ def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
         model_kwargs = conf_dict["module"]
         model = RegressionTask(**model_kwargs)
 
+        model.model = RegressionTestModel()
+
         # Instantiate trainer
         trainer = Trainer(fast_dev_run=True, log_every_n_steps=1)
         trainer.fit(model=model, datamodule=datamodule)

tests/trainers/test_segmentation.py

@@ -2,11 +2,14 @@
 # Licensed under the MIT License.
 
 import os
-from typing import Any, Dict, Type, cast
+from typing import Any, Dict, Generator, Type, cast
 
 import pytest
+import segmentation_models_pytorch as smp
+from _pytest.monkeypatch import MonkeyPatch
 from omegaconf import OmegaConf
 from pytorch_lightning import LightningDataModule, Trainer
+from torch.nn.modules import Module
 
 from torchgeo.datamodules import (
     ChesapeakeCVPRDataModule,
@@ -18,6 +21,12 @@
 )
 from torchgeo.trainers import SemanticSegmentationTask
 
+from .test_utils import SegmentationTestModel
+
+
+def create_model(**kwargs: Any) -> Module:
+    return SegmentationTestModel(**kwargs)
+
 
 class TestSemanticSegmentationTask:
     @pytest.mark.parametrize(
@@ -35,7 +44,12 @@ class TestSemanticSegmentationTask:
             ("sen12ms_s2_reduced", SEN12MSDataModule),
         ],
     )
-    def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
+    def test_trainer(
+        self,
+        monkeypatch: Generator[MonkeyPatch, None, None],
+        name: str,
+        classname: Type[LightningDataModule],
+    ) -> None:
         conf = OmegaConf.load(os.path.join("conf", "task_defaults", name + ".yaml"))
         conf_dict = OmegaConf.to_object(conf.experiment)
         conf_dict = cast(Dict[Any, Dict[Any, Any]], conf_dict)
@@ -45,6 +59,10 @@ def test_trainer(self, name: str, classname: Type[LightningDataModule]) -> None:
         datamodule = classname(**datamodule_kwargs)
 
         # Instantiate model
+        monkeypatch.setattr(smp, "Unet", create_model)  # type: ignore[attr-defined]
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            smp, "DeepLabV3Plus", create_model
+        )
         model_kwargs = conf_dict["module"]
         model = SemanticSegmentationTask(**model_kwargs)
 

tests/trainers/test_utils.py

@@ -3,6 +3,7 @@
 
 import os
 from pathlib import Path
+from typing import Any, cast
 
 import pytest
 import torch
@@ -16,6 +17,43 @@
 )
 
 
+class ClassificationTestModel(Module):
+    def __init__(
+        self, in_chans: int = 3, num_classes: int = 1000, **kwargs: Any
+    ) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(  # type: ignore[attr-defined]
+            in_channels=in_chans, out_channels=1, kernel_size=1
+        )
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))  # type: ignore[attr-defined]
+        self.fc = nn.Linear(1, num_classes)  # type: ignore[attr-defined]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        x = self.pool(x)
+        x = torch.flatten(x, 1)  # type: ignore[attr-defined]
+        x = self.fc(x)
+        return x
+
+
+class RegressionTestModel(ClassificationTestModel):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(in_chans=3, num_classes=1)
+
+
+class SegmentationTestModel(Module):
+    def __init__(
+        self, in_channels: int = 3, classes: int = 1000, **kwargs: Any
+    ) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(  # type: ignore[attr-defined]
+            in_channels=in_channels, out_channels=classes, kernel_size=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return cast(torch.Tensor, self.conv1(x))
+
+
 def test_extract_encoder_unsupported_model(tmp_path: Path) -> None:
     checkpoint = {"hyper_parameters": {"some_unsupported_model": "resnet18"}}
     path = os.path.join(str(tmp_path), "dummy.ckpt")