Optuna integration reports AttributeError

🐛 Bug

When using PyTorchLightningPruningCallback to search best hyperparams, it reports AttributeError: 'AcceleratorConnector' object has no attribute 'distributed_backend'

To Reproduce

from typing import List, Optional

import optuna
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torchmetrics
import torchvision
from optuna.integration.pytorch_lightning import PyTorchLightningPruningCallback
from torch.utils.data import random_split, DataLoader


class FashionDataModule(pl.LightningDataModule):
    def __init__(self, data_dir: str, batch_size: int):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def setup(self, stage: Optional[str] = None):
        self.train_set = torchvision.datasets.FashionMNIST(
            self.data_dir, train=True, download=True, transform=torchvision.transforms.ToTensor()
        )
        self.test_set = torchvision.datasets.FashionMNIST(
            self.data_dir, train=False, download=True, transform=torchvision.transforms.ToTensor()
        )
        self.train_set, self.valid_set = random_split(self.train_set, [55000, 5000])

    def train_dataloader(self) -> DataLoader:
        return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True, num_workers=4)

    def val_dataloader(self) -> DataLoader:
        return DataLoader(self.valid_set, batch_size=self.batch_size, shuffle=False, num_workers=4)

    def test_dataloader(self) -> DataLoader:
        return DataLoader(self.test_set, batch_size=self.batch_size, shuffle=False, num_workers=4)


class SimpleNet(nn.Module):
    def __init__(self, d_hids: List[int], p_drop: float):
        super(SimpleNet, self).__init__()

        hidden_layers = []
        d_inp = 28 * 28
        for d_hid in d_hids:
            hidden_layers.append(nn.Linear(d_inp, d_hid))
            hidden_layers.append(nn.ReLU())
            hidden_layers.append(nn.Dropout(p_drop))
            d_inp = d_hid
        hidden_layers.append(nn.Linear(d_inp, 10))

        self.layers = nn.Sequential(*hidden_layers)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.layers(inputs)


class LitSimpleNet(pl.LightningModule):
    def __init__(self, d_hids: List[int], p_drop: float):
        super().__init__()
        self.model = SimpleNet(d_hids, p_drop)
        self.criterion = nn.CrossEntropyLoss()
        self.accuracy = torchmetrics.Accuracy()

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.model(inputs.view(-1, 28 * 28))

    def training_step(self, batch, batch_idx) -> torch.Tensor:
        inputs, targets = batch
        outputs = self(inputs)
        return self.criterion(outputs, targets)

    def validation_step(self, batch, batch_idx):
        inputs, targets = batch
        outputs = self(inputs)
        self.accuracy(outputs, targets)
        self.log("valid_acc", self.accuracy, on_step=False, on_epoch=True, prog_bar=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=3e-4, weight_decay=1e-5)


def objective(trial: optuna.trial.Trial) -> float:
    n_layers = trial.suggest_int("n_layers", 1, 3)
    p_drop = trial.suggest_float("p_drop", 0.1, 0.5)
    d_hids = [trial.suggest_int(f"d_hid_{i}", 16, 128, log=True) for i in range(n_layers)]

    datamodule = FashionDataModule(".", 128)
    model = LitSimpleNet(d_hids, p_drop)
    trainer = pl.Trainer(
        max_epochs=20,
        accelerator="gpu",
        devices=1,
        enable_checkpointing=False,
        logger=True,
        default_root_dir=".",
        callbacks=[PyTorchLightningPruningCallback(trial, monitor="valid_acc")]
    )

    hparams = dict(n_layers=n_layers, d_hids=d_hids, p_drop=p_drop)
    trainer.logger.log_hyperparams(hparams)
    trainer.fit(model, datamodule=datamodule)
    return trainer.callback_metrics["valid_acc"].item()


if __name__ == "__main__":
    pruner = optuna.pruners.MedianPruner()
    study = optuna.create_study(direction="maximize", pruner=pruner)
    study.optimize(objective, n_trials=100, timeout=1000)

    print("Number of Finished Trials:", len(study.trials))

    trial = study.best_trial
    print("Best Trial:")
    print("\tValue:", trial.value)
    print("\tParams:")
    for key, value in trial.params.items():
        print(f"\t\t{key}: {value}")

[W 2022-09-08 20:14:45,294] Trial 0 failed because of the following error: AttributeError("'AcceleratorConnector' object has no attribute 'distributed_backend'")
Traceback (most recent call last):
  File "/home/wyn/miniconda3/envs/wyn/lib/python3.8/site-packages/optuna/study/_optimize.py", line 196, in _run_trial
    value_or_values = func(trial)
  File "optuna_examples/optuna_lightning_example.py", line 89, in objective
    trainer = pl.Trainer(
  File "/home/wyn/miniconda3/envs/wyn/lib/python3.8/site-packages/pytorch_lightning/utilities/argparse.py", line 345, in insert_env_defaults
    return fn(self, **kwargs)
  File "/home/wyn/miniconda3/envs/wyn/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 497, in __init__
    self._call_callback_hooks("on_init_start")
  File "/home/wyn/miniconda3/envs/wyn/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1585, in _call_callback_hooks
    fn(self, *args, **kwargs)
  File "/home/wyn/miniconda3/envs/wyn/lib/python3.8/site-packages/optuna/integration/pytorch_lightning.py", line 61, in on_init_start
    trainer._accelerator_connector.distributed_backend is not None  # type: ignore
AttributeError: 'AcceleratorConnector' object has no attribute 'distributed_backend'

Expected behavior

Should not report any errors.

Environment

• Lightning Component: Trainer
• PyTorch Lightning Version: 1.7.5
• PyTorch Version: 1.12.1
• Python version: 3.8.13
• OS: Linux (Ubuntu 20.04)
• CUDA/cuDNN version: 11.3.1
• How you installed PyTorch: conda

cc @akihironitta

[nzw0301]

@hrzn Hi, I’m from the Optuna-dev team. Optuna’s pytorch-lightning (PL) integration module doesn’t support PL>=1.6 because it broke backwards-compatibility as investigated in https://github.com/optuna/optuna/issues/3418. Unfortunately, Optuna team doesn’t have time to fix the module soon to support recent PL; we would like to wait for a PR from optuna and PL users.

@tchaton I believe you can close this issue because the issue comes from Optuna…

[mikiotada]

With Optuna==3.0.2 with lightning==1.5.10, I got ValueError: optuna.integration.PyTorchLightningPruningCallback supports only optuna.storages.RDBStorage in DDP. After downgrading Optuna to 2.0.0 (arbitrary version) while keeping lightning==1.5.10, it ran without any error.