Update model.eval based on the new lightning

Signed-off-by: Samet Akcay <[email protected]>

Author: samet-akcay

pyproject.toml

@@ -34,8 +34,8 @@ core = [
     "opencv-python>=4.5.3.56",
     "pandas>=1.1.0",
     "timm>=0.5.4,<=0.9.16",
-    "lightning>2,<2.2.0",
-    "torch>=2,<2.3.0",
+    "lightning>=2.2",
+    "torch>=2",
     "torchmetrics>=1.3.2",
     "open-clip-torch>=2.23.0",
 ]

src/anomalib/models/image/cflow/lightning_model.py

@@ -84,6 +84,7 @@ def __init__(
         # TODO(ashwinvaidya17): LR should be part of optimizer in config.yaml since  cflow has custom optimizer.
         # CVS-122670
         self.learning_rate = lr
+        self.model.encoder.eval()
 
     def configure_optimizers(self) -> Optimizer:
         """Configure optimizers for each decoder.
@@ -119,7 +120,6 @@ def training_step(self, batch: dict[str, str | torch.Tensor], *args, **kwargs) -
         del args, kwargs  # These variables are not used.
 
         opt = self.optimizers()
-        self.model.encoder.eval()
 
         images: torch.Tensor = batch["image"]
         activation = self.model.encoder(images)

src/anomalib/models/image/csflow/lightning_model.py

@@ -6,7 +6,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 import logging
 from typing import Any
 
@@ -68,6 +67,7 @@ def _setup(self) -> None:
             clamp=self.clamp,
             num_channels=self.num_channels,
         )
+        self.model.feature_extractor.eval()
 
     def training_step(self, batch: dict[str, str | torch.Tensor], *args, **kwargs) -> STEP_OUTPUT:
         """Perform the training step of CS-Flow.
@@ -82,7 +82,6 @@ def training_step(self, batch: dict[str, str | torch.Tensor], *args, **kwargs) -
         """
         del args, kwargs  # These variables are not used.
 
-        self.model.feature_extractor.eval()
         z_dist, jacobians = self.model(batch["image"])
         loss = self.loss(z_dist, jacobians)
         self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)

src/anomalib/models/image/fastflow/lightning_model.py

@@ -6,7 +6,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 from typing import Any
 
 import torch
@@ -52,9 +51,8 @@ def __init__(
         self.conv3x3_only = conv3x3_only
         self.hidden_ratio = hidden_ratio
 
-        self.loss = FastflowLoss()
-
         self.model: FastflowModel
+        self.loss = FastflowLoss()
 
     def _setup(self) -> None:
         if self.input_size is None:

src/anomalib/models/image/padim/lightning_model.py

@@ -6,7 +6,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 import logging
 
 import torch
@@ -52,7 +51,8 @@ def __init__(
             pre_trained=pre_trained,
             layers=layers,
             n_features=n_features,
-        ).eval()
+        )
+        self.model.feature_extractor.eval()
 
         self.stats: list[torch.Tensor] = []
         self.embeddings: list[torch.Tensor] = []
@@ -75,9 +75,7 @@ def training_step(self, batch: dict[str, str | torch.Tensor], *args, **kwargs) -
         """
         del args, kwargs  # These variables are not used.
 
-        self.model.feature_extractor.eval()
         embedding = self.model(batch["image"])
-
         self.embeddings.append(embedding.cpu())
 
     def fit(self) -> None:

src/anomalib/models/image/padim/torch_model.py

@@ -3,7 +3,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 from random import sample
 from typing import TYPE_CHECKING
 
@@ -67,8 +66,8 @@ class PadimModel(nn.Module):
 
     def __init__(
         self,
-        layers: list[str],
         backbone: str = "resnet18",
+        layers: list[str] = ["layer1", "layer2", "layer3"],  # noqa: B006
         pre_trained: bool = True,
         n_features: int | None = None,
     ) -> None:
@@ -77,7 +76,11 @@ def __init__(
 
         self.backbone = backbone
         self.layers = layers
-        self.feature_extractor = TimmFeatureExtractor(backbone=self.backbone, layers=layers, pre_trained=pre_trained)
+        self.feature_extractor = TimmFeatureExtractor(
+            backbone=self.backbone,
+            layers=layers,
+            pre_trained=pre_trained,
+        ).eval()
         self.n_features_original = sum(self.feature_extractor.out_dims)
         self.n_features = n_features or _N_FEATURES_DEFAULTS.get(self.backbone)
         if self.n_features is None:

src/anomalib/models/image/patchcore/lightning_model.py

@@ -78,9 +78,7 @@ def training_step(self, batch: dict[str, str | torch.Tensor], *args, **kwargs) -
         """
         del args, kwargs  # These variables are not used.
 
-        self.model.feature_extractor.eval()
         embedding = self.model(batch["image"])
-
         self.embeddings.append(embedding)
 
     def fit(self) -> None:

src/anomalib/models/image/patchcore/torch_model.py

@@ -49,7 +49,7 @@ def __init__(
             backbone=self.backbone,
             pre_trained=pre_trained,
             layers=self.layers,
-        )
+        ).eval()
         self.feature_pooler = torch.nn.AvgPool2d(3, 1, 1)
         self.anomaly_map_generator = AnomalyMapGenerator()
 

src/anomalib/models/image/reverse_distillation/lightning_model.py

@@ -6,7 +6,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 from collections.abc import Sequence
 from typing import Any
 
@@ -50,9 +49,8 @@ def __init__(
         self.layers = layers
         self.anomaly_map_mode = anomaly_map_mode
 
-        self.loss = ReverseDistillationLoss()
-
         self.model: ReverseDistillationModel
+        self.loss = ReverseDistillationLoss()
 
     def _setup(self) -> None:
         if self.input_size is None:

src/anomalib/models/image/stfpm/lightning_model.py

@@ -6,7 +6,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 from collections.abc import Sequence
 from typing import Any
 
@@ -61,7 +60,6 @@ def training_step(self, batch: dict[str, str | torch.Tensor], *args, **kwargs) -
         """
         del args, kwargs  # These variables are not used.
 
-        self.model.teacher_model.eval()
         teacher_features, student_features = self.model.forward(batch["image"])
         loss = self.loss(teacher_features, student_features)
         self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)

src/anomalib/models/image/stfpm/torch_model.py

@@ -3,7 +3,6 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
@@ -36,7 +35,7 @@ def __init__(
         self.tiler: Tiler | None = None
 
         self.backbone = backbone
-        self.teacher_model = TimmFeatureExtractor(backbone=self.backbone, pre_trained=True, layers=layers)
+        self.teacher_model = TimmFeatureExtractor(backbone=self.backbone, pre_trained=True, layers=layers).eval()
         self.student_model = TimmFeatureExtractor(
             backbone=self.backbone,
             pre_trained=False,

src/anomalib/models/image/uflow/feature_extraction.py

@@ -32,15 +32,13 @@ def get_feature_extractor(backbone: str, input_size: tuple[int, int] = (256, 256
         msg = f"Feature extractor must be one of {AVAILABLE_EXTRACTORS}."
         raise ValueError(msg)
 
+    feature_extractor: nn.Module
     if backbone in ["resnet18", "wide_resnet50_2"]:
-        return FeatureExtractor(backbone, input_size, layers=("layer1", "layer2", "layer3"))
+        feature_extractor = FeatureExtractor(backbone, input_size, layers=("layer1", "layer2", "layer3")).eval()
     if backbone == "mcait":
-        return MCaitFeatureExtractor()
-    msg = (
-        "`backbone` must be one of `[mcait, resnet18, wide_resnet50_2]`. These are the only feature extractors tested. "
-        "It does not mean that other feature extractors will not work."
-    )
-    raise ValueError(msg)
+        feature_extractor = MCaitFeatureExtractor().eval()
+
+    return feature_extractor
 
 
 class FeatureExtractor(TimmFeatureExtractor):

src/anomalib/models/image/uflow/torch_model.py

@@ -3,7 +3,6 @@
 # Copyright (C) 2023-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 import torch
 from FrEIA import framework as ff
 from FrEIA import modules as fm

tests/integration/model/test_models.py

@@ -6,7 +6,6 @@
 # Copyright (C) 2023-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-
 from pathlib import Path
 
 import pytest
@@ -144,12 +143,10 @@ def test_export(
             dataset_path (Path): Root to dataset from fixture.
             project_path (Path): Path to temporary project folder from fixture.
         """
-        if model_name == "reverse_distillation":
-            # TODO(ashwinvaidya17): Restore this test after fixing reverse distillation
+        if model_name in ("reverse_distillation", "rkde"):
+            # TODO(ashwinvaidya17): Restore this test after fixing the issue
             # https://github.com/openvinotoolkit/anomalib/issues/1513
-            pytest.skip("Reverse distillation fails to convert to ONNX")
-        elif model_name == "rkde" and export_type == ExportType.OPENVINO:
-            pytest.skip("RKDE fails to convert to OpenVINO")
+            pytest.skip(f"{model_name} fails to convert to ONNX and OpenVINO")
 
         model, dataset, engine = self._get_objects(
             model_name=model_name,