Early stopping with PyTorchLightningPruningCallback fails with OSError: [Errno 24] Too many open files

[hermda02]

Expected behavior

Expected pruning of runs to pass cleanly with a run fully stopped before moving onto the next.

Environment

• Optuna version: 3.5.0
• Python version: 3.12.1
• OS: Ubuntu 22.04.2
• PyTorch version: 2.2.0
• PyTorchLightning version: 2.2.0.post0
• mlflow version: 2.10.2

Error messages, stack traces, or logs

[I 2024-03-18 10:34:09,963] Trial 34 pruned. Trial was pruned at epoch 0.██████████████████████████████████| 48/48 [00:00<00:00, 124.59it/s]
================================================
------------ Hyper-parameter Tuning ------------
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

   | Name     | Type              | Params
------------------------------------------------
0  | loss_fn  | BCEWithLogitsLoss | 0
1  | inc      | Down              | 2.8 K
2  | down1    | Down              | 14.0 K
3  | down2    | Down              | 55.7 K
4  | down3    | Down              | 221 K
5  | down4    | Down              | 886 K
6  | up1      | Up                | 574 K
7  | up2      | Up                | 143 K
8  | up3      | Up                | 36.1 K
9  | up4      | Up                | 9.1 K
10 | outc     | OutConv           | 17
11 | accuracy | BinaryAccuracy    | 0
------------------------------------------------
1.9 M     Trainable params
0         Non-trainable params
1.9 M     Total params
7.776     Total estimated model params size (MB)
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [00:08<00:00, 50.98it/s]
[W 2024-03-18 10:34:22,046] Trial 35 failed with parameters: {'lr': 0.0006124846948934273, 'bs': 1, 'beta1': 0.6647434103439328, 'beta2': 0.
9469497665681813} because of the following error: OSError(24, 'Too many open files').
Traceback (most recent call last):
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/optuna/study/_optimize.py", line 200, in _run_trial
    value_or_values = func(trial)
                      ^^^^^^^^^^^
  File "/home/chronos/mainstay/mainstay_ml_scripts/ml_model/hyperparam_opt.py", line 68, in objective
    trainer.fit(model, datamodule=datamodule)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 543, in fit
    call._call_and_handle_interrupt(
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_han
dle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 579, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 986, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 1032, in _run_sta
ge
    self.fit_loop.run()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py", line 205, in run
    self.advance()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py", line 363, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 139, in
 run
    self.on_advance_end(data_fetcher)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 287, in
 on_advance_end
    self.val_loop.run()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/utilities.py", line 182, in _decorato
r
    return loop_run(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 113, in run
    self.reset()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 231, in res
et
    iter(data_fetcher)  # creates the iterator inside the fetcher
    ^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fetchers.py", line 104, in __iter__
    super().__iter__()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fetchers.py", line 51, in __iter__

    self.iterator = iter(self.combined_loader)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/utilities/combined_loader.py", line 351, in __iter__
    iter(iterator)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/utilities/combined_loader.py", line 155, in __iter__
    self._load_current_iterator()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/utilities/combined_loader.py", line 173, in _load_current_iterator
    self.iterators = [iter(self.iterables[self._iterator_idx])]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 439, in __iter__
    return self._get_iterator()
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 387, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1040, in __init__
    w.start()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
                  ^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/context.py", line 282, in _Popen
    return Popen(process_obj)
           ^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/popen_fork.py", line 65, in _launch
    child_r, parent_w = os.pipe()
                        ^^^^^^^^^
OSError: [Errno 24] Too many open files
[W 2024-03-18 10:34:22,052] Trial 35 failed with value None.
Traceback (most recent call last):
  File "/home/chronos/mainstay/mainstay_ml_scripts/ml_model/train.py", line 78, in <module>
    run_hp_optimizer(study_name="mainstay-optuna", n_trials=50)
  File "/home/chronos/mainstay/mainstay_ml_scripts/ml_model/hyperparam_opt.py", line 87, in run_hp_optimizer
    study.optimize(objective, n_trials=n_trials)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/optuna/study/study.py", line 451, in optimize
    _optimize(
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/optuna/study/_optimize.py", line 66, in _optimize
    _optimize_sequential(
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/optuna/study/_optimize.py", line 163, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/optuna/study/_optimize.py", line 251, in _run_trial
    raise func_err
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/optuna/study/_optimize.py", line 200, in _run_trial
    value_or_values = func(trial)
                      ^^^^^^^^^^^
  File "/home/chronos/mainstay/mainstay_ml_scripts/ml_model/hyperparam_opt.py", line 68, in objective
trainer.fit(model, datamodule=datamodule)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 543, in fit
    call._call_and_handle_interrupt(
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 579, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 986, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 1032, in _run_stage
    self.fit_loop.run()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py", line 205, in run
    self.advance()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py", line 363, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 139, in run
    self.on_advance_end(data_fetcher)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 287, in on_advance_end
    self.val_loop.run()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/utilities.py", line 182, in _decorator
    return loop_run(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 113, in run
    self.reset()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 231, in reset
    iter(data_fetcher)  # creates the iterator inside the fetcher
    ^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fetchers.py", line 104, in __iter__
    super().__iter__()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/loops/fetchers.py", line 51, in __iter__
    self.iterator = iter(self.combined_loader)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/utilities/combined_loader.py", line 351, in __iter__
    iter(iterator)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/utilities/combined_loader.py", line 155, in __iter__
    self._load_current_iterator()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/lightning/pytorch/utilities/combined_loader.py", line 173, in _load_current_iterator
    self.iterators = [iter(self.iterables[self._iterator_idx])]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 439, in __iter__
    return self._get_iterator()
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 387, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1040, in __init__
    w.start()
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
                  ^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/context.py", line 282, in _Popen
    return Popen(process_obj)
           ^^^^^^^^^^^^^^^^^^
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/home/chronos/miniconda3/envs/mainstay_ml/lib/python3.12/multiprocessing/popen_fork.py", line 65, in _launch
    child_r, parent_w = os.pipe()
                        ^^^^^^^^^
OSError: [Errno 24] Too many open files
Epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [04:02<00:00,  1.78it/s]
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [03:49<00:00,  1.88it/s]
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 216/216 [03:36<00:00,  1.00it/s]
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 108/108 [01:42<00:00,  1.06it/s]
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [01:30<00:00,  4.75it/s]
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 216/216 [01:18<00:00,  2.76it/s]
Epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [00:57<00:00,  7.55it/s]
Epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 216/216 [00:35<00:00,  6.15it/s]
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [00:22<00:00, 18.99it/s]
Epoch 0: 100%|██████████| 432/432 [00:10<00:00, 42.19it/s]

Steps to reproduce

def objective(trial):
    print("================================================")
    print("------------ Hyper-parameter Tuning ------------")
    metrics_callback = MetricsCallback()
    early_stop_callback=PyTorchLightningPruningCallback(trial, monitor="val_loss")
    trainer = Trainer(
        logger=False,
        max_epochs=config.num_epochs,
        callbacks=[metrics_callback, early_stop_callback],
        devices=[1],
    )
    best_val_loss = float('Inf')

    num_workers = 20
    PIN_MEMORY = True
    PERSISTENT_WORKERS = True
    imagePaths = sorted(list(paths.list_images(config.IMG_DATASET_PATH)))[:config.NUM_SAMPLES]
    maskPaths = sorted(list(paths.list_images(config.MASK_DATASET_PATH)))[:config.NUM_SAMPLES]


    with mlflow.start_run():
        lr, bs, betas = suggest_hyperparameters(trial)
        batch_size = 2**bs

        datamodule = ImageDataModule(imagePaths, maskPaths, batch_size, num_workers, PIN_MEMORY, PERSISTENT_WORKERS)
       
        config.log_config_params()
        mlflow.log_params(trial.params)
        mlflow.log_param('device', config.device)

        # Initialize model
        model = UNet(in_layers=config.IMG_DEPTH, n_classes=int(config.N_CLASSES), learning_rate=lr, betas=betas, kernel_size=config.kernel_size)
        trainer.fit(model, datamodule=datamodule)

        if metrics_callback.metrics[-1]["val_loss"].item() < best_val_loss:
            best_val_loss = metrics_callback.metrics[-1]["val_loss"].item()

    return metrics_callback.metrics[-1]["val_loss"]


# python code

Additional context (optional)

Results seem to be similar to those found on the PyTorch forums: https://discuss.pytorch.org/t/too-many-open-files-caused-by-persistent-workers-and-pin-memory/193372, though in my case is not dependent on the pin_memory or persistent_workers variables.

Watching the memory usage on our compute node indicates a similar behavior -- more and more files remain unclosed after pruning until eventually the Python environment throws an error saying too many files are open.

This does not fail upon the first pruned run, though an error begins to appear after several runs have been pruned. Normally a run takes on the order of 8 seconds per epoch, after some time it appears that several runs have stacked up with some individual epochs taking on the order of 5 minutes to complete:

`================================================
------------ Hyper-parameter Tuning ------------
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

| Name | Type | Params

0 | loss_fn | BCEWithLogitsLoss | 0
1 | inc | Down | 2.8 K
2 | down1 | Down | 14.0 K
3 | down2 | Down | 55.7 K
4 | down3 | Down | 221 K
5 | down4 | Down | 886 K
6 | up1 | Up | 574 K
7 | up2 | Up | 143 K
8 | up3 | Up | 36.1 K
9 | up4 | Up | 9.1 K
10 | outc | OutConv | 17
11 | accuracy | BinaryAccuracy | 0

1.9 M Trainable params
0 Non-trainable params
1.9 M Total params
7.776 Total estimated model params size (MB)
Epoch 0: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 27/27 [07:58<00:00, 0.06it/s$
Epoch 0: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 27/27 [08:36<00:00, 0.05it/s$
Epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 108/108 [08:16<00:00, 0.22it/s$
Epoch 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [11:24<00:00, 0.63it/s$
Epoch 4: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 216/216 [09:45<00:00, 0.37it/s$
Epoch 4: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [09:37<00:00, 0.75it/s$
Epoch 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [00:09<00:00, 47.91it/s$
[I 2024-03-18 10:30:30,473] Trial 24 pruned. Trial was pruned at epoch 1.██████████████████████████████████| 48/48 [00:00<00:00, 133.81it/s$

`

[nzw0301]

Hi, could you share minimal reproducible codes with us?

[nzw0301]

I suspect ImageDataModule open files but not close them.

[hermda02]

Unfortunately this is the shortest version I was able to reproduce in a reasonable time:

from torchvision import tv_tensors
from torchvision.io import read_image
from torch.utils.data import Dataset
from torchvision.transforms import v2
from imutils import paths
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from lightning import LightningDataModule

from torchmetrics import Accuracy
import torch.nn as nn
import lightning as L
import torch.nn.functional as F 

from typing import Optional

from lightning import LightningModule, Trainer
from lightning.pytorch import Callback
import optuna
from optuna.integration.mlflow import MLflowCallback
from optuna.integration import PyTorchLightningPruningCallback
import mlflow


class ImageDataModule(LightningDataModule):
    def __init__(self, image_paths: str, mask_paths: str, batch_size: int, num_workers: int,
                 pin_memory: bool, persistent_workers: bool):
        super().__init__()
        self.persistent_workers = persistent_workers
        self.pin_memory = pin_memory
        self.num_workers = num_workers
        self.batch_size = batch_size
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = v2.Compose([
            v2.Resize((240, 320), antialias=True),
            v2.RandomVerticalFlip(p=0.5),
            v2.RandomHorizontalFlip(p=0.5),
        ])
        self.valid_split = 0.1
        self.test_split = 0.1

    def setup(self, stage: Optional[str]) -> None:
        if stage == "fit":
            split = train_test_split(self.image_paths, self.mask_paths, test_size=self.valid_split)

            (self.train_images, self.valid_images) = split[:2]
            (self.train_masks, self.valid_masks) = split[2:]

        if stage == "test":
            total_split = self.test_split + self.valid_split
            ratio = self.test_split/total_split
            split = train_test_split(self.image_paths, self.mask_paths, test_size=total_split)

            (self.train_images, split_images) = split[:2]
            (self.train_masks, split_masks) = split[2:]

            split2 = train_test_split(split_images, split_masks, test_size=ratio)

            (self.test_images, self.valid_images) = split2[:2]
            (self.test_masks, self.valid_masks) = split2[2:]

    def train_dataloader(self) -> DataLoader:
        return DataLoader(dataset=ImageData(self.train_images, self.train_masks, transform=self.transform),
                          shuffle=True, batch_size=self.batch_size, num_workers=self.num_workers)
    
    def val_dataloader(self) -> DataLoader:
        return DataLoader(dataset=ImageData(self.valid_images, self.valid_masks, transform=self.transform),
                          shuffle=False, batch_size=self.batch_size, num_workers=self.num_workers)

    def test_dataloader(self) -> DataLoader:
        return DataLoader(dataset=ImageData(self.test_images, self.test_masks, transform=self.transform),
                          shuffle=False, batch_size=self.batch_size, num_workers=self.num_workers)

class ImageData(Dataset):
    def __init__(self, image_paths, mask_paths, transform=None):
        """
        Arguments:
            csv_file (string): Path to csv file with annotations.
            img_dir (string): Image directory.
            transform (callable, optional): Optional transformations to be 
                applied to the sample.
        """
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        # Find the image and mask location
        image_path = self.image_paths[idx]
        mask_path = self.mask_paths[idx]

        image = read_image(image_path)
        mask = tv_tensors.Mask(read_image(mask_path))#.bool()

        if self.transform is not None:
            image_out, mask_out = self.transform(image, mask)
        else:
            image_out, mask_out = image, mask

        sample = {'image': image_out, 'mask': mask_out, 'path': image_path}
        return sample
    

class Down(L.LightningModule):
    def __init__(self, in_layers, out_layers, dropout_prob=0, max_pooling=True):
        super().__init__()
        self.conv = nn.Sequential()
        if max_pooling:
            self.conv.add_module("maxpool_1", nn.MaxPool2d(2))
        self.conv.add_module("conv_1", nn.Conv2d(in_layers, out_layers, kernel_size=3, padding='same'))
        self.conv.add_module("batchnorm1", nn.BatchNorm2d(out_layers))
        self.conv.add_module("relu_1", nn.ReLU())
        self.conv.add_module("conv_2", nn.Conv2d(out_layers, out_layers, kernel_size=3, padding='same'))
        self.conv.add_module("batchnorm2", nn.BatchNorm2d(out_layers))
        self.conv.add_module("relu_2", nn.ReLU())
        if dropout_prob > 0:
            self.conv.add_module("dropout", nn.Dropout(dropout_prob))

    def forward(self, x):
        return self.conv(x)
    
class Up(L.LightningModule):
    def __init__(self, in_layers, out_layers):
        super().__init__()

        self.up = nn.ConvTranspose2d(in_layers, in_layers // 2, kernel_size=2, stride=2)

        self.conv = nn.Sequential()
        self.conv.add_module("conv_1", nn.Conv2d(in_layers, out_layers, kernel_size=3, padding='same'))
        self.conv.add_module("batchnorm1", nn.BatchNorm2d(out_layers))
        self.conv.add_module("relu_1", nn.ReLU())
        self.conv.add_module("conv_2", nn.Conv2d(out_layers, out_layers, kernel_size=3, padding='same'))
        self.conv.add_module("batchnorm2", nn.BatchNorm2d(out_layers))
        self.conv.add_module("relu_2", nn.ReLU())

    def forward(self, x1, x2):
        # ConvTranspose the previous layer
        x1 = self.up(x1)

        # Ensure sizes of x1 and x2 agree
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        # Add together the old and the new
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)
    
class OutConv(L.LightningModule):
    def __init__(self, in_layers, out_layers):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_layers, out_layers, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

class UNet(L.LightningModule):
    def __init__(self, in_layers, num_classes, learning_rate, betas=(0.9,0.999), kernel_size=2):
        super(UNet, self).__init__()
        self.in_layers = in_layers
        self.num_classes = num_classes
        self.kernel_size = kernel_size
        self.lr = learning_rate
        self.betas = betas

        self.loss_fn = nn.BCEWithLogitsLoss()

        self.inc = (Down(in_layers, 16, max_pooling=False))
        self.down1 = (Down(16, 32))
        self.down2 = (Down(32, 64))
        self.down3 = (Down(64, 128))
        self.down4 = (Down(128, 256))
        self.up1 = (Up(256, 128))
        self.up2 = (Up(128, 64))
        self.up3 = (Up(64, 32))
        self.up4 = (Up(32, 16))
        self.outc = (OutConv(16, self.num_classes))
        if self.num_classes > 1:
            self.accuracy = Accuracy("multiclass", num_classes=self.num_classes)
        else:
            self.accuracy = Accuracy("binary")

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits
    
    def training_step(self, batch, batch_idx):
        x, y = batch['image'].float(), batch['mask'].float()
        batch_size = len(x)
        logits = self.forward(x)
        loss = self.loss_fn(logits, y)
        # preds = torch.argmax(logits, dim=1)
        acc = self.accuracy(logits, y)

        self.log("train_loss", loss, on_epoch=True, batch_size=batch_size)
        self.log("train_acc", acc, on_epoch=True, batch_size=batch_size)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch['image'].float(), batch['mask'].float()
        batch_size = len(x)
        logits = self.forward(x)
        val_loss = self.loss_fn(logits, y)
        val_acc = self.accuracy(logits, y)
        self.log("val_loss", val_loss, on_epoch=True, batch_size=batch_size)
        self.log("val_acc", val_acc, on_epoch=True, batch_size=batch_size)
        return {'val_loss' : val_loss, 'val_acc' : val_acc}
    
    def test_step(self, batch, batch_idx):
        x, y = batch['image'].float(), batch['mask'].float()
        y_hat = self.forward(x)
        test_acc = self.accuracy(y_hat, y)
        test_loss = self.loss_fn(y_hat, y)

        self.log("test_loss", test_loss)
        self.log("test_acc", test_acc)
        return test_loss
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr, betas=self.betas)


IMG_DATASET_PATH = '/path/to/img_data/'
MASK_DATASET_PATH = '/path/to/mask_data/'

num_workers = 20
PIN_MEMORY = False
PERSISTENT_WORKERS = False
imagePaths = sorted(list(paths.list_images(IMG_DATASET_PATH)))[:960]
maskPaths = sorted(list(paths.list_images(MASK_DATASET_PATH)))[:960]

datamodule = ImageDataModule(imagePaths, maskPaths, 8, num_workers, PIN_MEMORY, PERSISTENT_WORKERS)

model = UNet(in_layers=3, num_classes=1, learning_rate=1.5e-3, kernel_size=3)


class MetricsCallback(Callback):
    def __init__(self):
        super().__init__()
        self.metrics = []

    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        self.metrics.append(trainer.callback_metrics)
        # return super().on_validation_end(trainer, pl_module)

def objective(trial):
    print("================================================")
    print("------------ Hyper-parameter Tuning ------------")
    
    metrics_callback = MetricsCallback()
    early_stop_callback=PyTorchLightningPruningCallback(trial, monitor="val_loss")
    trainer = Trainer(
        logger=False,
        max_epochs=5,
        callbacks=[metrics_callback, early_stop_callback],
        devices=[1],
    )
    best_val_loss = float('Inf')

    num_workers = 20
    PIN_MEMORY = True
    PERSISTENT_WORKERS = True
    imagePaths = sorted(list(paths.list_images(IMG_DATASET_PATH)))[:500]
    maskPaths = sorted(list(paths.list_images(MASK_DATASET_PATH)))[:500]

    with mlflow.start_run():
        lr, bs, betas = suggest_hyperparameters(trial)
        batch_size = 2**bs

        # Load in the data
        datamodule = ImageDataModule(imagePaths, maskPaths, 8, num_workers, PIN_MEMORY, PERSISTENT_WORKERS)

       
        mlflow.log_params(trial.params)
        mlflow.log_param('device', 'cuda')

        # Initialize model
        model = UNet(in_layers=3, num_classes=1, learning_rate=1.5e-3, kernel_size=3)
        trainer.fit(model, datamodule=datamodule)

        if metrics_callback.metrics[-1]["val_loss"].item() < best_val_loss:
            best_val_loss = metrics_callback.metrics[-1]["val_loss"].item()

    return metrics_callback.metrics[-1]["val_loss"]


def suggest_hyperparameters(trial):
    lr = trial.suggest_float("lr", 1e-4, 1e-1, log=True)
    bs = trial.suggest_int("bs", 1, 6)
    beta1 = trial.suggest_float("beta1", 0.5, 0.9999)
    beta2 = trial.suggest_float("beta2", 0.9, 0.9999)
    betas = (beta1, beta2)
    # optimizer_name = trial.suggest_categorical("optimizer_name", ["Adam", "Adadelta"])
    return lr, bs, betas

def run_hp_optimizer(study_name: str, n_trials: int):
    study = optuna.create_study(study_name=study_name, direction="minimize")
    study.optimize(objective, n_trials=n_trials)

    # Print optuna study statistics
    print("\n++++++++++++++++++++++++++++++++++\n")
    print("Study statistics: ")
    print("  Number of finished trials: ", len(study.trials))

    print("Best trial:")
    trial = study.best_trial

    print("  Trial number: ", trial.number)
    print("  Loss (trial value): ", trial.value)

    print("  Params: ")
    for key, value in trial.params.items():
        print("    {}: {}".format(key, value))


run_hp_optimizer(study_name="mlflow-optuna", n_trials=50)```

[nzw0301]

Thanks and sorry for the late response. I'm not sure but pytorch/vision#539 seems related, so maybe not using read_image is a possible workaround.

[nzw0301]

I think the issue looks more question rather than the bug of optuna, let me covert this issue to the discussion.