Use sklearn in runif (Lightning-AI#15426)

* Use sklearn in runif
* test by removing sklearn dep
* remove repeated code
* seed

Author: rohitgr7

tests/legacy/simple_classif_training.py

@@ -14,141 +14,16 @@
 import os
 
 import torch
-import torch.nn.functional as F
-from sklearn.datasets import make_classification
-from sklearn.model_selection import train_test_split
-from torch import nn
-from torch.utils.data import DataLoader, Dataset
-from torchmetrics import Accuracy
 
 import pytorch_lightning as pl
-from pytorch_lightning import LightningDataModule, LightningModule, seed_everything
+from pytorch_lightning import seed_everything
 from pytorch_lightning.callbacks import EarlyStopping
+from tests_pytorch.helpers.datamodules import ClassifDataModule
+from tests_pytorch.helpers.simple_models import ClassificationModel
 
 PATH_LEGACY = os.path.dirname(__file__)
 
 
-class SklearnDataset(Dataset):
-    def __init__(self, x, y, x_type, y_type):
-        self.x = x
-        self.y = y
-        self._x_type = x_type
-        self._y_type = y_type
-
-    def __getitem__(self, idx):
-        return torch.tensor(self.x[idx], dtype=self._x_type), torch.tensor(self.y[idx], dtype=self._y_type)
-
-    def __len__(self):
-        return len(self.y)
-
-
-class SklearnDataModule(LightningDataModule):
-    def __init__(self, sklearn_dataset, x_type, y_type, batch_size: int = 128):
-        super().__init__()
-        self.batch_size = batch_size
-        self._x, self._y = sklearn_dataset
-        self._split_data()
-        self._x_type = x_type
-        self._y_type = y_type
-
-    def _split_data(self):
-        self.x_train, self.x_test, self.y_train, self.y_test = train_test_split(
-            self._x, self._y, test_size=0.20, random_state=42
-        )
-        self.x_train, self.x_predict, self.y_train, self.y_predict = train_test_split(
-            self._x, self._y, test_size=0.20, random_state=42
-        )
-        self.x_train, self.x_valid, self.y_train, self.y_valid = train_test_split(
-            self.x_train, self.y_train, test_size=0.40, random_state=42
-        )
-
-    def train_dataloader(self):
-        return DataLoader(
-            SklearnDataset(self.x_train, self.y_train, self._x_type, self._y_type),
-            shuffle=True,
-            batch_size=self.batch_size,
-        )
-
-    def val_dataloader(self):
-        return DataLoader(
-            SklearnDataset(self.x_valid, self.y_valid, self._x_type, self._y_type), batch_size=self.batch_size
-        )
-
-    def test_dataloader(self):
-        return DataLoader(
-            SklearnDataset(self.x_test, self.y_test, self._x_type, self._y_type), batch_size=self.batch_size
-        )
-
-    def predict_dataloader(self):
-        return DataLoader(
-            SklearnDataset(self.x_predict, self.y_predict, self._x_type, self._y_type), batch_size=self.batch_size
-        )
-
-
-class ClassifDataModule(SklearnDataModule):
-    def __init__(self, num_features=24, length=6000, num_classes=3, batch_size=128):
-        data = make_classification(
-            n_samples=length,
-            n_features=num_features,
-            n_classes=num_classes,
-            n_clusters_per_class=2,
-            n_informative=int(num_features / num_classes),
-            random_state=42,
-        )
-        super().__init__(data, x_type=torch.float32, y_type=torch.long, batch_size=batch_size)
-
-
-class ClassificationModel(LightningModule):
-    def __init__(self, num_features=24, num_classes=3, lr=0.01):
-        super().__init__()
-        self.save_hyperparameters()
-
-        self.lr = lr
-        for i in range(3):
-            setattr(self, f"layer_{i}", nn.Linear(num_features, num_features))
-            setattr(self, f"layer_{i}a", torch.nn.ReLU())
-        setattr(self, "layer_end", nn.Linear(num_features, num_classes))
-
-        self.train_acc = Accuracy()
-        self.valid_acc = Accuracy()
-        self.test_acc = Accuracy()
-
-    def forward(self, x):
-        x = self.layer_0(x)
-        x = self.layer_0a(x)
-        x = self.layer_1(x)
-        x = self.layer_1a(x)
-        x = self.layer_2(x)
-        x = self.layer_2a(x)
-        x = self.layer_end(x)
-        logits = F.softmax(x, dim=1)
-        return logits
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
-        return [optimizer], []
-
-    def training_step(self, batch, batch_idx):
-        x, y = batch
-        logits = self.forward(x)
-        loss = F.cross_entropy(logits, y)
-        self.log("train_loss", loss, prog_bar=True)
-        self.log("train_acc", self.train_acc(logits, y), prog_bar=True)
-        return {"loss": loss}
-
-    def validation_step(self, batch, batch_idx):
-        x, y = batch
-        logits = self.forward(x)
-        self.log("val_loss", F.cross_entropy(logits, y), prog_bar=False)
-        self.log("val_acc", self.valid_acc(logits, y), prog_bar=True)
-
-    def test_step(self, batch, batch_idx):
-        x, y = batch
-        logits = self.forward(x)
-        self.log("test_loss", F.cross_entropy(logits, y), prog_bar=False)
-        self.log("test_acc", self.test_acc(logits, y), prog_bar=True)
-
-
 def main_train(dir_path, max_epochs: int = 20):
     seed_everything(42)
     stopping = EarlyStopping(monitor="val_acc", mode="max", min_delta=0.005)

tests/tests_pytorch/accelerators/test_hpu.py

@@ -76,8 +76,8 @@ def test_all_stages(tmpdir, hpus):
     trainer.predict(model)
 
 
-@RunIf(hpu=True)
-@mock.patch.dict(os.environ, os.environ.copy())
+@RunIf(hpu=True, sklearn=True)
+@mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 def test_optimization(tmpdir):
     seed_everything(42)
 

tests/tests_pytorch/accelerators/test_ipu.py

@@ -149,7 +149,8 @@ def test_inference_only(tmpdir, devices):
     trainer.predict(model)
 
 
-@RunIf(ipu=True)
+@RunIf(ipu=True, sklearn=True)
+@mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 def test_optimization(tmpdir):
     seed_everything(42)
 

tests/tests_pytorch/callbacks/test_early_stopping.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import logging
 import math
+import os
 import pickle
 from typing import List, Optional
 from unittest import mock
@@ -56,6 +57,8 @@ def on_train_epoch_end(self, trainer, pl_module):
         self.saved_states.append(self.state_dict().copy())
 
 
+@RunIf(sklearn=True)
+@mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 def test_resume_early_stopping_from_checkpoint(tmpdir):
     """Prevent regressions to bugs:
 
@@ -98,6 +101,7 @@ def test_resume_early_stopping_from_checkpoint(tmpdir):
         new_trainer.fit(model, datamodule=dm, ckpt_path=checkpoint_filepath)
 
 
+@RunIf(sklearn=True)
 def test_early_stopping_no_extraneous_invocations(tmpdir):
     """Test to ensure that callback methods aren't being invoked outside of the callback handler."""
     model = ClassificationModel()
@@ -195,6 +199,7 @@ def test_pickling(tmpdir):
     assert vars(early_stopping) == vars(early_stopping_loaded)
 
 
+@RunIf(sklearn=True)
 def test_early_stopping_no_val_step(tmpdir):
     """Test that early stopping callback falls back to training metrics when no validation defined."""
 

tests/tests_pytorch/callbacks/test_lr_monitor.py

@@ -22,6 +22,7 @@
 from pytorch_lightning.demos.boring_classes import BoringModel
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests_pytorch.helpers.datamodules import ClassifDataModule
+from tests_pytorch.helpers.runif import RunIf
 from tests_pytorch.helpers.simple_models import ClassificationModel
 
 
@@ -284,6 +285,7 @@ def configure_optimizers(self):
     assert all(len(lr) == expected_number_logged for lr in lr_monitor.lrs.values())
 
 
+@RunIf(sklearn=True)
 def test_lr_monitor_param_groups(tmpdir):
     """Test that learning rates are extracted and logged for single lr scheduler."""
 

tests/tests_pytorch/callbacks/test_quantization.py

@@ -35,7 +35,7 @@
 @pytest.mark.parametrize("observe", ["average", "histogram"])
 @pytest.mark.parametrize("fuse", [True, False])
 @pytest.mark.parametrize("convert", [True, False])
-@RunIf(quantization=True, max_torch="1.11")
+@RunIf(quantization=True, sklearn=True, max_torch="1.11")
 def test_quantization(tmpdir, observe: str, fuse: bool, convert: bool):
     """Parity test for quant model."""
     cuda_available = CUDAAccelerator.is_available()
@@ -100,7 +100,7 @@ def test_quantization(tmpdir, observe: str, fuse: bool, convert: bool):
     assert torch.allclose(org_score, quant2_score, atol=0.45)
 
 
-@RunIf(quantization=True)
+@RunIf(quantization=True, sklearn=True)
 def test_quantize_torchscript(tmpdir):
     """Test converting to torchscipt."""
     dm = RegressDataModule()
@@ -116,7 +116,7 @@ def test_quantize_torchscript(tmpdir):
     tsmodel(tsmodel.quant(batch[0]))
 
 
-@RunIf(quantization=True)
+@RunIf(quantization=True, sklearn=True)
 def test_quantization_exceptions(tmpdir):
     """Test wrong fuse layers."""
     with pytest.raises(MisconfigurationException, match="Unsupported qconfig"):
@@ -157,7 +157,7 @@ def custom_trigger_last(trainer):
     "trigger_fn,expected_count",
     [(None, 9), (3, 3), (custom_trigger_never, 0), (custom_trigger_even, 5), (custom_trigger_last, 2)],
 )
-@RunIf(quantization=True)
+@RunIf(quantization=True, sklearn=True)
 def test_quantization_triggers(tmpdir, trigger_fn: Union[None, int, Callable], expected_count: int):
     """Test  how many times the quant is called."""
     dm = RegressDataModule()
@@ -216,7 +216,7 @@ def test_quantization_disable_observers(tmpdir, observer_enabled_stages):
         )
 
 
-@RunIf(quantization=True)
+@RunIf(quantization=True, sklearn=True)
 def test_quantization_val_test_predict(tmpdir):
     """Test the default quantization aware training not affected by validating, testing and predicting."""
     seed_everything(42)

tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py

@@ -23,6 +23,9 @@
 import pytorch_lightning as pl
 from pytorch_lightning import Callback, Trainer
 from tests_pytorch import _PATH_LEGACY
+from tests_pytorch.helpers.datamodules import ClassifDataModule
+from tests_pytorch.helpers.runif import RunIf
+from tests_pytorch.helpers.simple_models import ClassificationModel
 
 LEGACY_CHECKPOINTS_PATH = os.path.join(_PATH_LEGACY, "checkpoints")
 CHECKPOINT_EXTENSION = ".ckpt"
@@ -32,18 +35,17 @@
 
 
 @pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
+@RunIf(sklearn=True)
 def test_load_legacy_checkpoints(tmpdir, pl_version: str):
     PATH_LEGACY = os.path.join(LEGACY_CHECKPOINTS_PATH, pl_version)
     with patch("sys.path", [PATH_LEGACY] + sys.path):
-        from simple_classif_training import ClassifDataModule, ClassificationModel
-
         path_ckpts = sorted(glob.glob(os.path.join(PATH_LEGACY, f"*{CHECKPOINT_EXTENSION}")))
         assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
         path_ckpt = path_ckpts[-1]
 
-        model = ClassificationModel.load_from_checkpoint(path_ckpt)
+        model = ClassificationModel.load_from_checkpoint(path_ckpt, num_features=24)
         trainer = Trainer(default_root_dir=str(tmpdir))
-        dm = ClassifDataModule()
+        dm = ClassifDataModule(num_features=24, length=6000, batch_size=128, n_clusters_per_class=2, n_informative=8)
         res = trainer.test(model, datamodule=dm)
         assert res[0]["test_loss"] <= 0.7
         assert res[0]["test_acc"] >= 0.85
@@ -62,6 +64,7 @@ def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningMo
 
 
 @pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
+@RunIf(sklearn=True)
 def test_legacy_ckpt_threading(tmpdir, pl_version: str):
     def load_model():
         import torch
@@ -84,17 +87,16 @@ def load_model():
 
 
 @pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
+@RunIf(sklearn=True)
 def test_resume_legacy_checkpoints(tmpdir, pl_version: str):
     PATH_LEGACY = os.path.join(LEGACY_CHECKPOINTS_PATH, pl_version)
     with patch("sys.path", [PATH_LEGACY] + sys.path):
-        from simple_classif_training import ClassifDataModule, ClassificationModel
-
         path_ckpts = sorted(glob.glob(os.path.join(PATH_LEGACY, f"*{CHECKPOINT_EXTENSION}")))
         assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
         path_ckpt = path_ckpts[-1]
 
-        dm = ClassifDataModule()
-        model = ClassificationModel()
+        dm = ClassifDataModule(num_features=24, length=6000, batch_size=128, n_clusters_per_class=2, n_informative=8)
+        model = ClassificationModel(num_features=24)
         stop = LimitNbEpochs(1)
 
         trainer = Trainer(

tests/tests_pytorch/core/test_datamodules.py

@@ -148,6 +148,7 @@ def test_dm_pickle_after_init():
     pickle.dumps(dm)
 
 
+@RunIf(sklearn=True)
 def test_train_loop_only(tmpdir):
     seed_everything(7)
 
@@ -169,6 +170,7 @@ def test_train_loop_only(tmpdir):
     assert trainer.callback_metrics["train_loss"] < 1.1
 
 
+@RunIf(sklearn=True)
 def test_train_val_loop_only(tmpdir):
     seed_everything(7)
 
@@ -226,6 +228,7 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         assert dm.my_state_dict == {"my": "state_dict"}
 
 
+@RunIf(sklearn=True)
 def test_full_loop(tmpdir):
     seed_everything(7)