🐞 NotImplementedError: ('{} cannot be pickled', '_SingleProcessDataLoaderIter')

[deyantara1303]

Describe the bug

NotImplemented Error exception raised when using pre-processor module while creating model.
It gets raised after first epoch during save checkpoint stage.

Dataset

Custom

Model

EfficientAD

Steps to reproduce the behavior

Steps to reproduce the behaviour:

1. Install anomalib v2.0.0.b3
2. Create train.py with the following:
3. run train.py

`from anomalib.engine import Engine
from anomalib.models import EfficientAd
from anomalib.data import Folder

if name == "main":

datamodule = Folder(
    name="dummy",
    root="datasets/dummy",
    normal_dir="train/good",
    abnormal_dir="val/anomaly",
    normal_test_dir= "val/good",
    # mask_dir: "ground_truth/xxx"
    normal_split_ratio= 0,
    extensions= [".jpg"],
    train_batch_size= 1,
    eval_batch_size= 1,
    num_workers= 0,
    val_split_mode= "same_as_test"
)

pre_processor = EfficientAd.configure_pre_processor(image_size=(512, 512))
model = EfficientAd(pre_processor=pre_processor)
engine = Engine(max_epochs=10)

engine.fit(datamodule=datamodule, model=model)`

OS information

OS information:

• OS: MacOS
• Python version: 3.10.15
• Anomalib version: 2.0.0b3
• PyTorch version: 2.8.0
• CUDA/cuDNN version: MPS
• GPU models and configuration: Apple Silicon M3
• Any other relevant information: I am using custom dataset

Also, tried with CUDA and Nvidia GPU, facing same error.

Expected behavior

Training should continue to next epoch after saving model.ckpt.

Screenshots

No response

Pip/GitHub

pip

What version/branch did you use?

24.2

Configuration YAML

hparams.yaml 

After passing preprocessor, this looks like:
--------------------------------------------
evaluator: true
imagenet_dir: ./datasets/imagenette
lr: 0.0001
model_size: S
pad_maps: true
padding: false
post_processor: true
pre_processor: !!python/object:anomalib.pre_processing.pre_processing.PreProcessor
  _backward_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _backward_pre_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _buffers: {}
  _forward_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _forward_hooks_always_called: !!python/object/apply:collections.OrderedDict
  - []
  _forward_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
  - []
  _forward_pre_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _forward_pre_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
  - []
  _is_full_backward_hook: null
  _load_state_dict_post_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _load_state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _modules:
    export_transform: !!python/object:torchvision.transforms.v2._container.Compose
      _backward_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _backward_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _buffers: {}
      _forward_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _forward_hooks_always_called: !!python/object/apply:collections.OrderedDict
      - []
      _forward_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
      - []
      _forward_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _forward_pre_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
      - []
      _is_full_backward_hook: null
      _load_state_dict_post_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _load_state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _modules: {}
      _non_persistent_buffers_set: !!set {}
      _parameters: {}
      _state_dict_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      training: true
      transforms:
      - !!python/object:torchvision.transforms.v2._geometry.Resize
        _backward_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _backward_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _buffers: {}
        _forward_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _forward_hooks_always_called: !!python/object/apply:collections.OrderedDict
        - []
        _forward_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
        - []
        _forward_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _forward_pre_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
        - []
        _is_full_backward_hook: null
        _load_state_dict_post_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _load_state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _modules: {}
        _non_persistent_buffers_set: !!set {}
        _parameters: {}
        _state_dict_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        antialias: false
        interpolation: &id001 !!python/object/apply:torchvision.transforms.functional.InterpolationMode
        - bilinear
        max_size: null
        size:
        - 512
        - 512
        training: true
    transform: !!python/object:torchvision.transforms.v2._container.Compose
      _backward_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _backward_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _buffers: {}
      _forward_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _forward_hooks_always_called: !!python/object/apply:collections.OrderedDict
      - []
      _forward_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
      - []
      _forward_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _forward_pre_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
      - []
      _is_full_backward_hook: null
      _load_state_dict_post_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _load_state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _modules: {}
      _non_persistent_buffers_set: !!set {}
      _parameters: {}
      _state_dict_hooks: !!python/object/apply:collections.OrderedDict
      - []
      _state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
      - []
      training: true
      transforms:
      - !!python/object:torchvision.transforms.v2._geometry.Resize
        _backward_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _backward_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _buffers: {}
        _forward_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _forward_hooks_always_called: !!python/object/apply:collections.OrderedDict
        - []
        _forward_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
        - []
        _forward_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _forward_pre_hooks_with_kwargs: !!python/object/apply:collections.OrderedDict
        - []
        _is_full_backward_hook: null
        _load_state_dict_post_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _load_state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _modules: {}
        _non_persistent_buffers_set: !!set {}
        _parameters: {}
        _state_dict_hooks: !!python/object/apply:collections.OrderedDict
        - []
        _state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
        - []
        antialias: true
        interpolation: *id001
        max_size: null
        size:
        - 512
        - 512
        training: true
  _non_persistent_buffers_set: !!set {}
  _parameters: {}
  _state_dict_hooks: !!python/object/apply:collections.OrderedDict
  - []
  _state_dict_pre_hooks: !!python/object/apply:collections.OrderedDict
  - []
  training: true
teacher_out_channels: 384
visualizer: true
weight_decay: 1.0e-05


The one that works, without passing pre-processor:
---------------------------------------------------
pre_processor: true
post_processor: true
evaluator: true
visualizer: true
imagenet_dir: ./datasets/imagenette
teacher_out_channels: 384
model_size: S
lr: 0.0001
weight_decay: 1.0e-05
padding: false
pad_maps: true

Logs

Validation DataLoader 0: 100%|██████████| 5/5 [00:01<00:00,  4.40it/s]
Epoch 0: 100%|██████████| 14/14 [00:13<00:00,  1.02it/s, train_st_step=9.390, train_ae_step=0.801, train_stae_step=0.00346, train_loss_step=10.20]Memory consumed val_end: 0.08 GB 

Epoch 0: 100%|██████████| 14/14 [00:13<00:00,  1.02it/s, train_st_step=9.390, train_ae_step=0.801, train_stae_step=0.00346, train_loss_step=10.20, train_st_epoch=13.50, train_ae_epoch=0.965, train_stae_epoch=0.00315, train_loss_epoch=14.50]Traceback (most recent call last):
  File "/Users/deya/PycharmProjects/Experiments_AD/anomaly_detection_anomalib/run_train_v2.py", line 36, in <module>
    engine.fit(datamodule=datamodule, model=model)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/anomalib/engine/engine.py", line 416, in fit
    self.trainer.fit(model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 538, in fit
    call._call_and_handle_interrupt(
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 47, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 574, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 981, in _run
    results = self._run_stage()
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1025, in _run_stage
    self.fit_loop.run()
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py", line 206, in run
    self.on_advance_end()
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py", line 378, in on_advance_end
    call._call_callback_hooks(trainer, "on_train_epoch_end", monitoring_callbacks=True)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 218, in _call_callback_hooks
    fn(trainer, trainer.lightning_module, *args, **kwargs)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 325, in on_train_epoch_end
    self._save_topk_checkpoint(trainer, monitor_candidates)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 387, in _save_topk_checkpoint
    self._save_none_monitor_checkpoint(trainer, monitor_candidates)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 715, in _save_none_monitor_checkpoint
    self._save_checkpoint(trainer, filepath)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 390, in _save_checkpoint
    trainer.save_checkpoint(filepath, self.save_weights_only)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1365, in save_checkpoint
    self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py", line 490, in save_checkpoint
    self.checkpoint_io.save_checkpoint(checkpoint, filepath, storage_options=storage_options)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/fabric/plugins/io/torch_io.py", line 58, in save_checkpoint
    _atomic_save(checkpoint, path)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/lightning/fabric/utilities/cloud_io.py", line 85, in _atomic_save
    torch.save(checkpoint, bytesbuffer)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/torch/serialization.py", line 965, in save
    _save(
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/torch/serialization.py", line 1211, in _save
    pickler.dump(obj)
  File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 762, in __getstate__
    raise NotImplementedError("{} cannot be pickled", self.__class__.__name__)
NotImplementedError: ('{} cannot be pickled', '_SingleProcessDataLoaderIter')

Code of Conduct

• I agree to follow this project's Code of Conduct

[alexriedel1]

Yes I'm also experiencing this error when passing a pre processor to efficientad. Im still investigating why this is happening.

[deyantara1303]

Is there any other way to perform size augmentation? I tried adding train_augmentations during the data module creation using Folder, but the changes did not apply.

[samet-akcay]

hmm, is this only happening when using Folder, or with any dataset?

[deyantara1303]

if __name__ == main:  
 datamodule = MVTec(
        root="./datasets/MVTec",  # Path to download/store the dataset
        category="bottle",  # MVTec category to use
        train_batch_size=1,  # Number of images per training batch
        eval_batch_size=1,  # Number of images per validation/test batch
        num_workers=0,  # Number of parallel processes for data loading
    )

    logger = AnomalibTensorBoardLogger(save_dir="tb_logs")
    pre_processor = EfficientAd.configure_pre_processor(image_size=(512, 512))
    model = EfficientAd(pre_processor=pre_processor)
    engine = Engine(max_epochs=10, logger=logger)

    engine.fit(datamodule=datamodule, model=model)

With MVTEC dataset, I get the following error on 1st epoch end,

fn(trainer, trainer.lightning_module, *args, **kwargs)

File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/anomalib/post_processing/one_class.py", line 121, in on_validation_epoch_end
self._pixel_threshold = self._pixel_threshold_metric.compute()
File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/torchmetrics/metric.py", line 633, in wrapped_func
value = _squeeze_if_scalar(compute(*args, **kwargs))
File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/anomalib/metrics/base.py", line 191, in compute
return super().compute() # type: ignore[misc]
File "/Users/deya/miniconda3/envs/py310/lib/python3.10/site-packages/anomalib/metrics/threshold/f1_adaptive_threshold.py", line 98, in compute
return thresholds if thresholds.dim() == 0 else thresholds[torch.argmax(f1_score)]
IndexError: index 12475498 is out of bounds for dimension 0 with size 12475498

Without pre-processor, it works as expected.

[PepperoniPingu]

I'm no expert in the anomalib architecture but this is what I've found out:

When the EfficientAd class is initialized it calls HyperparametersMixin.save_hyperparameters() here . Since this is done before the pre processor is configured it normally does not configure the pre processor to be saved and no problems arise. But when the pre processor is explicitly configured as an argument (i.e. model = EfficientAd(pre_processor=...)), save_hyperparameters() recognizes the pre processor as a component that needs to be saved. Thus when it comes time to save a checkpoint it fails. Why it can't save the pre processor, I don't know.

A quick solution that works is to to set the pre processor after initialization:

model = EfficientAd()
model.pre_processor = EfficientAd.configure_pre_processor(IMAGE_SIZE)

Though probably not the correct way to fix this.

[samet-akcay]

I'm having some trouble finding the root cause of this problem. So far, I only managed to reproduce this issue on efficient-ad only.

There is a relevant issue created on PyTorch Lightning repo, which tries to pickle the models directly. I'm not sure if we are doing it anywhere though.
Lightning-AI/pytorch-lightning#17124

[markstanl]

Not sure if this is the exact same issue. I get the same behavior as already described.
However, I get the same error, with a similar-looking traceback when I run the following after I have properly trained a model on MVTec AD without the preprocessor.

engine.export(model, ExportType.Torch)

Results:

756 def __getstate__(self):
    757     # TODO: add limited pickling support for sharing an iterator
    758     # across multiple threads for HOGWILD.
    759     # Probably the best way to do this is by moving the sample pushing
    760     # to a separate thread and then just sharing the data queue
    761     # but signalling the end is tricky without a non-blocking API
--> [762](file:///C:/Users/MarkStanley/coding/pcb-anomaly-detection/.venv/Lib/site-packages/torch/utils/data/dataloader.py:762)     raise NotImplementedError("{} cannot be pickled", self.__class__.__name__)

NotImplementedError: ('{} cannot be pickled', '_SingleProcessDataLoaderIter')

[samet-akcay]

@markstanl, can you share your code so I can reproduce the exact same error?

[markstanl]

@markstanl, can you share your code so I can reproduce the exact same error?

It's the same code as discussed before.

datamodule = MVTecAD(
      root="./datasets/MVTecAD",
      category="bottle",
      train_batch_size=1,
      eval_batch_size=1,
      num_workers=8,
)
model = EfficientAd()
engine = Engine(max_epochs=2)
engine.fit(datamodule=datamodule, model=model)

This works just as expected. But adding the export causes the crash

datamodule = MVTecAD(
      root="./datasets/MVTecAD",
      category="bottle",
      train_batch_size=1,
      eval_batch_size=1,
      num_workers=8,
)
model = EfficientAd()
engine = Engine(max_epochs=2)
engine.fit(datamodule=datamodule, model=model)
engine.export(model, ExportType.Torch)

Results:

  File "anomaly.py", line 112, in <module>
    engine.export(model, ExportType.TORCH)
    ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\anomalib\engine\engine.py", line 802, in export
    exported_model_path = model.to_torch(
        export_root=export_root,
        model_file_name=model_file_name,
    )
  File "\.venv\Lib\site-packages\anomalib\models\components\base\export_mixin.py", line 105, in to_torch
    torch.save(
    ~~~~~~~~~~^
        obj={"model": self},
        ^^^^^^^^^^^^^^^^^^^^
        f=pt_model_path,
        ^^^^^^^^^^^^^^^^
    )
    ^
  File ".venv\Lib\site-packages\torch\serialization.py", line 965, in save
    _save(
    ~~~~~^
        obj,
        ^^^^
    ...<3 lines>...
        _disable_byteorder_record,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "\.venv\Lib\site-packages\torch\serialization.py", line 1211, in _save
    pickler.dump(obj)
    ~~~~~~~~~~~~^^^^^
  File ".venv\Lib\site-packages\torch\utils\data\dataloader.py", line 762, in __getstate__
    raise NotImplementedError("{} cannot be pickled", self.__class__.__name__)
NotImplementedError: ('{} cannot be pickled', '_SingleProcessDataLoaderIter')