Repository: optuna/optuna

Issue:
`optuna.integration.PyTorchLightningPruningCallback` does not work with `pytorch_lightning==0.7.1`.
## Environment

- Optuna version: 1.2.0
- Python version: 3.7.2
- OS: macOS 10.13.6
- (Optional) Other libraries and their versions: pytorch_lightning==0.7.1

## Error messages, stack traces, or logs

Command
```console
pytest tests/integration_tests/test_pytorch_lightning.py
```

Error message:
```console
self = <optuna.integration.pytorch_lightning.PyTorchLightningPruningCallback object at 0x7f2194069a20>
logs = {'accuracy': 0.0}

    def check_metrics(self, logs):
        monitor_val = logs.get(self.monitor)
        error_msg = (f'Early stopping conditioned on metric `{self.monitor}`'
                     f' which is not available. Available metrics are:'
                     f' `{"`, `".join(list(logs.keys()))}`')
    
        if monitor_val is None:
            if self.strict:
>               raise RuntimeError(error_msg)
E               RuntimeError: Early stopping conditioned on metric `val_loss` which is not available. Available metrics are: `accuracy`

venv/lib/python3.6/site-packages/pytorch_lightning/callbacks/early_stopping.py:81: RuntimeError

```

## Additional context

The base class of `PyTorchLightningPruningCallback` (i.e., `EarlyStopping`) is updated by 0.7.x.
For example,
1. the `check_metrics` method is newly added, and
1. the arguments of `on_epoch_end` are changed.

For further details, please refer to the following code:
- https://github.com/PyTorchLightning/pytorch-lightning/blob/0.7.1/pytorch_lightning/callbacks/early_stopping.py
- https://github.com/PyTorchLightning/pytorch-lightning/blob/0.5.3.3/pytorch_lightning/callbacks/pt_callbacks.py

To resolve this issue, we may need to consider backward compatibility. Do we support both v0.7.x and 0.5.x?


Discussion:

Resolving patch:
diff --git a/examples/pytorch_lightning_simple.py b/examples/pytorch_lightning_simple.py
--- a/examples/pytorch_lightning_simple.py
+++ b/examples/pytorch_lightning_simple.py
@@ -23,7 +23,7 @@
 import shutil
 
 import pytorch_lightning as pl
-from pytorch_lightning.logging import LightningLoggerBase
+from pytorch_lightning import Callback
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -35,8 +35,8 @@
 import optuna
 from optuna.integration import PyTorchLightningPruningCallback
 
-if pkg_resources.parse_version(pl.__version__) < pkg_resources.parse_version("0.6.0"):
-    raise RuntimeError("PyTorch Lightning>=0.6.0 is required for this example.")
+if pkg_resources.parse_version(pl.__version__) < pkg_resources.parse_version("0.7.1"):
+    raise RuntimeError("PyTorch Lightning>=0.7.1 is required for this example.")
 
 PERCENT_VALID_EXAMPLES = 0.1
 BATCHSIZE = 128
@@ -46,20 +46,15 @@
 MODEL_DIR = os.path.join(DIR, "result")
 
 
-class DictLogger(LightningLoggerBase):
-    """PyTorch Lightning `dict` logger."""
+class MetricsCallback(Callback):
+    """PyTorch Lightning metric callback."""
 
-    def __init__(self, version):
-        super(DictLogger, self).__init__()
+    def __init__(self):
+        super().__init__()
         self.metrics = []
-        self._version = version
 
-    def log_metrics(self, metric, step=None):
-        self.metrics.append(metric)
-
-    @property
-    def version(self):
-        return self._version
+    def on_validation_end(self, trainer, pl_module):
+        self.metrics.append(trainer.callback_metrics)
 
 
 class Net(nn.Module):
@@ -68,7 +63,7 @@ def __init__(self, trial):
         self.layers = []
         self.dropouts = []
 
-        # We optimize the number of layers, hidden untis in each layer and drouputs.
+        # We optimize the number of layers, hidden units in each layer and dropouts.
         n_layers = trial.suggest_int("n_layers", 1, 3)
         dropout = trial.suggest_uniform("dropout", 0.2, 0.5)
         input_dim = 28 * 28
@@ -99,18 +94,15 @@ def forward(self, data):
 class LightningNet(pl.LightningModule):
     def __init__(self, trial):
         super(LightningNet, self).__init__()
-
-        # Be careful not to overwrite `pl.LightningModule` attributes such as `self.model`.
-        self._model = Net(trial)
+        self.model = Net(trial)
 
     def forward(self, data):
-        return self._model(data)
+        return self.model(data)
 
     def training_step(self, batch, batch_nb):
         data, target = batch
         output = self.forward(data)
-        loss = F.nll_loss(output, target)
-        return {"loss": loss}
+        return {"loss": F.nll_loss(output, target)}
 
     def validation_step(self, batch, batch_nb):
         data, target = batch
@@ -118,17 +110,16 @@ def validation_step(self, batch, batch_nb):
         pred = output.argmax(dim=1, keepdim=True)
         correct = pred.eq(target.view_as(pred)).sum().item()
         accuracy = correct / data.size(0)
-        return {"validation_accuracy": accuracy}
+        return {"batch_val_acc": accuracy}
 
-    def validation_end(self, outputs):
-        accuracy = sum(x["validation_accuracy"] for x in outputs) / len(outputs)
+    def validation_epoch_end(self, outputs):
+        accuracy = sum(x["batch_val_acc"] for x in outputs) / len(outputs)
         # Pass the accuracy to the `DictLogger` via the `'log'` key.
-        return {"log": {"accuracy": accuracy}}
+        return {"log": {"val_acc": accuracy}}
 
     def configure_optimizers(self):
-        return Adam(self._model.parameters())
+        return Adam(self.model.parameters())
 
-    @pl.data_loader
     def train_dataloader(self):
         return torch.utils.data.DataLoader(
             datasets.MNIST(DIR, train=True, download=True, transform=transforms.ToTensor()),
@@ -136,7 +127,6 @@ def train_dataloader(self):
             shuffle=True,
         )
 
-    @pl.data_loader
     def val_dataloader(self):
         return torch.utils.data.DataLoader(
             datasets.MNIST(DIR, train=False, download=True, transform=transforms.ToTensor()),
@@ -146,31 +136,29 @@ def val_dataloader(self):
 
 
 def objective(trial):
-    # PyTorch Lightning will try to restore model parameters from previous trials if checkpoint
-    # filenames match. Therefore, the filenames for each trial must be made unique.
+    # Filenames for each trial must be made unique in order to access each checkpoint.
     checkpoint_callback = pl.callbacks.ModelCheckpoint(
-        os.path.join(MODEL_DIR, "trial_{}".format(trial.number)), monitor="accuracy"
+        os.path.join(MODEL_DIR, "trial_{}".format(trial.number), "{epoch}"), monitor="val_acc"
     )
 
     # The default logger in PyTorch Lightning writes to event files to be consumed by
-    # TensorBoard. We create a simple logger instead that holds the log in memory so that the
-    # final accuracy can be obtained after optimization. When using the default logger, the
-    # final accuracy could be stored in an attribute of the `Trainer` instead.
-    logger = DictLogger(trial.number)
-
+    # TensorBoard. We don't use any logger here as it requires us to implement several abstract
+    # methods. Instead we setup a simple callback, that saves metrics from each validation step.
+    metrics_callback = MetricsCallback()
     trainer = pl.Trainer(
-        logger=logger,
+        logger=False,
         val_percent_check=PERCENT_VALID_EXAMPLES,
         checkpoint_callback=checkpoint_callback,
         max_epochs=EPOCHS,
         gpus=0 if torch.cuda.is_available() else None,
-        early_stop_callback=PyTorchLightningPruningCallback(trial, monitor="accuracy"),
+        callbacks=[metrics_callback],
+        early_stop_callback=PyTorchLightningPruningCallback(trial, monitor="val_acc"),
     )
 
     model = LightningNet(trial)
     trainer.fit(model)
 
-    return logger.metrics[-1]["accuracy"]
+    return metrics_callback.metrics[-1]["val_acc"]
 
 
 if __name__ == "__main__":
diff --git a/optuna/integration/pytorch_lightning.py b/optuna/integration/pytorch_lightning.py
--- a/optuna/integration/pytorch_lightning.py
+++ b/optuna/integration/pytorch_lightning.py
@@ -6,6 +6,8 @@
 
 try:
     from pytorch_lightning.callbacks import EarlyStopping
+    from pytorch_lightning import LightningModule
+    from pytorch_lightning import Trainer
 
     _available = True
 except (ImportError, SyntaxError) as e:
@@ -39,17 +41,17 @@ class PyTorchLightningPruningCallback(EarlyStopping):
     def __init__(self, trial, monitor):
         # type: (optuna.trial.Trial, str) -> None
 
-        super(PyTorchLightningPruningCallback, self).__init__()
+        super(PyTorchLightningPruningCallback, self).__init__(monitor=monitor)
 
         _check_pytorch_lightning_availability()
 
         self._trial = trial
         self._monitor = monitor
 
-    def on_epoch_end(self, epoch, logs=None):
-        # type: (int, Optional[Dict[str, float]]) -> None
-
-        logs = logs or {}
+    def on_epoch_end(self, trainer, pl_module):
+        # type: (Trainer, LightningModule) -> None
+        logs = trainer.callback_metrics
+        epoch = pl_module.current_epoch
         current_score = logs.get(self._monitor)
         if current_score is None:
             return
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -50,8 +50,8 @@ def get_tests_require() -> List[str]:
 def get_extras_require() -> Dict[str, List[str]]:
 
     requirements = {
-        "checking": ["black", "hacking", "mypy",],
-        "codecov": ["codecov", "pytest-cov",],
+        "checking": ["black", "hacking", "mypy"],
+        "codecov": ["codecov", "pytest-cov"],
         "doctest": [
             "cma",
             "pandas",
@@ -60,7 +60,7 @@ def get_extras_require() -> Dict[str, List[str]]:
             "scikit-optimize",
             "mlflow",
         ],
-        "document": ["sphinx", "sphinx_rtd_theme",],
+        "document": ["sphinx", "sphinx_rtd_theme"],
         "example": [
             "catboost",
             "chainer",
@@ -86,9 +86,7 @@ def get_extras_require() -> Dict[str, List[str]]:
                 "dask[dataframe]",
                 "dask-ml",
                 "keras",
-                # TODO(toshihikoyanase): Remove the version constraint after resolving the issue
-                # https://github.com/optuna/optuna/issues/997.
-                "pytorch-lightning<0.7.0",
+                "pytorch-lightning>=0.7.1",
                 "tensorflow>=2.0.0",
                 "tensorflow-datasets",
             ]
@@ -118,19 +116,13 @@ def get_extras_require() -> Dict[str, List[str]]:
             "torchvision==0.5.0+cpu",
             "xgboost",
         ]
-        + (["allennlp<1", "fastai<2"] if (3, 5) < sys.version_info[:2] < (3, 8) else [])
-        + (["fastai<2"] if (3, 5) < sys.version_info[:2] < (3, 8) else [])
         + (
-            [
-                "keras",
-                # TODO(toshihikoyanase): Remove the version constraint after resolving the issue
-                # https://github.com/optuna/optuna/issues/997.
-                "pytorch-lightning<0.7.0",
-                "tensorflow",
-                "tensorflow-datasets",
-            ]
-            if sys.version_info[:2] < (3, 8)
+            ["allennlp<1", "fastai<2", "pytorch-lightning>=0.7.1"]
+            if (3, 5) < sys.version_info[:2] < (3, 8)
             else []
+        )
+        + (
+            ["keras", "tensorflow", "tensorflow-datasets"] if sys.version_info[:2] < (3, 8) else []
         ),
     }
 


Test patch: