🐞 fix(model): EfficientAD NotImplementedError: ('{} cannot be pickled', '_SingleProcessDataLoaderIter') (#2837)

* fix efficientad example and pickling error

* pre-commit hooks

Signed-off-by: Alexander Riedel <alex.riedel@googlemail.com>

---------

Signed-off-by: Alexander Riedel <alex.riedel@googlemail.com>

Author: alexriedel1

examples/api/03_models/efficient_ad.py

@@ -7,28 +7,38 @@
 particularly well-suited for industrial inspection tasks.
 """
 
+from torchvision.transforms.v2 import Compose, Resize
+
 from anomalib.data import MVTecAD
 from anomalib.engine import Engine
 from anomalib.models import EfficientAd
+from anomalib.pre_processing import PreProcessor
 
 # 1. Basic Usage
 # Initialize with default settings
 model = EfficientAd()
 
-# 2. Custom Configuration
+# OPTIONAL 2. Set up a pre-processing transformation
+transform = Compose([
+    Resize(size=(512, 512)),
+])
+pre_processor = PreProcessor(transform=transform)
+
+# 3. Custom Configuration
 # Configure model parameters
 model = EfficientAd(
     teacher_out_channels=384,  # Number of teacher output channels
-    model_size="m",
+    model_size="medium",
     lr=1e-4,
+    pre_processor=pre_processor,
 )
 
-# 3. Training Pipeline
+# 4. Training Pipeline
 # Set up the complete training pipeline
 datamodule = MVTecAD(
     root="./datasets/MVTecAD",
     category="bottle",
-    train_batch_size=32,
+    train_batch_size=1,
 )
 
 # Initialize training engine with specific settings
@@ -38,8 +48,9 @@
     devices=1,  # Number of devices to use
 )
 
-# Train the model
-engine.fit(
-    model=model,
-    datamodule=datamodule,
-)
+if __name__ == "__main__":
+    # Train the model
+    engine.fit(
+        model=model,
+        datamodule=datamodule,
+    )

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

@@ -380,6 +380,16 @@ def on_train_start(self) -> None:
             channel_mean_std = self.teacher_channel_mean_std(self.trainer.datamodule.train_dataloader())
             self.model.mean_std.update(channel_mean_std)
 
+    def __getstate__(self) -> dict:
+        """Modifies the python objects __getstate__ method.
+
+        To ensure that the imagenet iterator instance will not get pickled
+        when the model is saved, it needs to be removed from the objects dict.
+        """
+        state = self.__dict__.copy()
+        state.pop("imagenet_iterator", None)
+        return state
+
     def training_step(self, batch: Batch, *args, **kwargs) -> dict[str, torch.Tensor]:
         """Perform training step.