Repository: mlflow/mlflow

Issue:
[BUG] mlflow.autolog() for PyTorch Lightning logs `_step` metrics for epochs instead of steps
### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [x] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: Reproducible with example from [official docs](https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html)
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Docker image: `nvidia/cuda:10.2-base-ubuntu16.04`
- **MLflow installed from (source or binary)**: Python package
- **MLflow version (run ``mlflow --version``)**: 1.15.0
- **Python version**: 3.8.8
- **npm version, if running the dev UI**: No
- **Exact command to reproduce**: See code snippet

### Describe the problem
Expected behavior:
If I'm using `mlflow.autolog()` with a `pl.LightningModule` and use `self.log("train_loss", loss, on_epoch=True, on_step=True, logger=True)` inside its `training_step` function, MLflow should log `train_loss_step` metric for every logged **step** and `train_loss_epoch` for every logged **epoch**

Actual behavior:
`train_loss_step` and `train_loss_epoch` are recorded for every logged **epoch** only.

After running code in below sample, `train_loss_epoch`  contains 1 entry per epoch:
```
1617996945494 1.3506861925125122 0
1617996966448 1.201533317565918 1
1617996986735 1.1901590824127197 2
1617997007487 1.190608024597168 3
1617997028088 1.1838353872299194 4
1617997048012 1.1807078123092651 5
1617997068702 1.18239426612854 6
1617997089547 1.1766492128372192 7
1617997110542 1.175412654876709 8
1617997130747 1.1763886213302612 9
1617997150977 1.1739486455917358 10
1617997171487 1.1722205877304077 11
1617997191321 1.170091152191162 12
1617997211502 1.1701279878616333 13
1617997232425 1.1694077253341675 14
1617997252996 1.1664137840270996 15
1617997273641 1.1662592887878418 16
1617997294056 1.1691796779632568 17
1617997314113 1.164090871810913 18
1617997335002 1.1653193235397339 19
```
`train_loss_step` also contains 1 entry per **epoch**
```
1617996945494 1.0488566160202026 0
1617996966448 1.0117261409759521 1
1617996986735 0.9437647461891174 2
1617997007487 1.1345815658569336 3
1617997028088 0.9712427854537964 4
1617997048012 0.9615302681922913 5
1617997068702 0.9383407235145569 6
1617997089547 0.9366243481636047 7
1617997110542 0.9363885521888733 8
1617997130747 0.9403294324874878 9
1617997150977 0.9360908269882202 10
1617997171487 1.0088082551956177 11
1617997191321 0.9398048520088196 12
1617997211502 0.9855148792266846 13
1617997232425 0.9356205463409424 14
1617997252996 0.9355267286300659 15
1617997273641 0.937127947807312 16
1617997294056 0.9390522837638855 17
1617997314113 0.9355039000511169 18
1617997335002 0.9383139610290527 19
```


### Code to reproduce issue
Small modifications to [code snippet in documentation](https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html)

```
import os

import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

import mlflow.pytorch


class MNISTModel(pl.LightningModule):
    def __init__(self):
        super(MNISTModel, self).__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)

        # Use the current of PyTorch logger
        self.log("train_loss", loss, on_epoch=True, on_step=True, logger=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


def print_auto_logged_info(r):
    print("metrics: {}".format(r.data.metrics))


def main() -> None:
    mlflow.set_experiment("Autolog issue")

    # Initialize our model
    mnist_model = MNISTModel()

    # Initialize DataLoader from MNIST Dataset
    train_ds = MNIST(os.getcwd(), train=True,
                     download=True, transform=transforms.ToTensor())
    train_loader = DataLoader(train_ds, batch_size=32)

    # Initialize a trainer
    trainer = pl.Trainer(max_epochs=20, progress_bar_refresh_rate=20)

    # Auto log all MLflow entities
    mlflow.pytorch.autolog()

    # Train the model
    with mlflow.start_run() as run:
        trainer.fit(mnist_model, train_loader)

        # fetch the auto logged parameters and metrics
        print_auto_logged_info(mlflow.get_run(run_id=run.info.run_id))


if __name__ == "__main__":
    main()
```


### Other info / logs
PyTorch Lightning adds the `_step` and `_epoch` suffixes to the logged metrics, so on our side we used this info and slightly adjusted [__MLflowPLCallback](https://github.com/mlflow/mlflow/blob/86a03e9488f7cf2c7c0e46531fd2168923c79ec2/mlflow/pytorch/_pytorch_autolog.py#L83) by adding the following hook:
```
        def on_train_batch_end(
            self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
        ):
           # ADD CODE TO FILTER ON `_step` SUFFIX

           # Then we log filtered metrics on pl_module.global_step instead of  pl_module.current_epoch
           metrics_logger.record_metrics(step_metrics, pl_module.global_step)
```
We also adjusted `__MLflowPLCallback._log_epoch_metrics()` to filter out metrics with the `_step` suffix.

Are there any plans to make classes such as [__MLflowPLCallback](https://github.com/mlflow/mlflow/blob/86a03e9488f7cf2c7c0e46531fd2168923c79ec2/mlflow/pytorch/_pytorch_autolog.py#L83) accessible so that we can subclass them to add/modify such functionality?

Thanks!

### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [x] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
@danieltremblay Thanks for filing this. 
cc: @shrinath-suresh @HamidShojanazeri @chauhang Any comment on this as needed to filter out "_prefixs" to the autolog metrics?

Resolving patch:
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
--- a/mlflow/pytorch/__init__.py
+++ b/mlflow/pytorch/__init__.py
@@ -876,6 +876,7 @@ def load_state_dict(state_dict_uri, **kwargs):
 @autologging_integration(FLAVOR_NAME)
 def autolog(
     log_every_n_epoch=1,
+    log_every_n_step=None,
     log_models=True,
     disable=False,
     exclusive=False,
@@ -905,6 +906,9 @@ def autolog(
 
     :param log_every_n_epoch: If specified, logs metrics once every `n` epochs. By default, metrics
                        are logged after every epoch.
+    :param log_every_n_step: If specified, logs batch metrics once every `n` global step.
+                       By default, metrics are not logged for steps. Note that setting this to 1 can
+                       cause performance issues and is not recommended.
     :param log_models: If ``True``, trained models are logged as MLflow model artifacts.
                        If ``False``, trained models are not logged.
     :param disable: If ``True``, disables the PyTorch Lightning autologging integration.
diff --git a/mlflow/pytorch/_pytorch_autolog.py b/mlflow/pytorch/_pytorch_autolog.py
--- a/mlflow/pytorch/_pytorch_autolog.py
+++ b/mlflow/pytorch/_pytorch_autolog.py
@@ -8,6 +8,7 @@
 from pytorch_lightning.core.memory import ModelSummary
 from pytorch_lightning.utilities import rank_zero_only
 
+from mlflow.exceptions import MlflowException
 from mlflow.utils.autologging_utils import (
     ExceptionSafeAbstractClass,
     BatchMetricsLogger,
@@ -30,6 +31,9 @@
 # once the above mentioned issues have been addressed
 
 
+_pl_version = Version(pl.__version__)
+
+
 def _get_optimizer_name(optimizer):
     """
     In pytorch-lightining 1.1.0, `LightningOptimizer` was introduced:
@@ -57,15 +61,25 @@ class __MLflowPLCallback(pl.Callback, metaclass=ExceptionSafeAbstractClass):
     Callback for auto-logging metrics and parameters.
     """
 
-    def __init__(self, client, metrics_logger, run_id, log_models, log_every_n_epoch):
+    def __init__(
+        self, client, metrics_logger, run_id, log_models, log_every_n_epoch, log_every_n_step
+    ):
+        if log_every_n_step and _pl_version < Version("1.1.0"):
+            raise MlflowException(
+                "log_every_n_step is only supported for PyTorch-Lightning >= 1.1.0"
+            )
         self.early_stopping = False
         self.client = client
         self.metrics_logger = metrics_logger
         self.run_id = run_id
         self.log_models = log_models
         self.log_every_n_epoch = log_every_n_epoch
+        self.log_every_n_step = log_every_n_step
+        # Sets for tracking which metrics are logged on steps and which are logged on epochs
+        self._step_metrics = set()
+        self._epoch_metrics = set()
 
-    def _log_metrics(self, trainer, pl_module):
+    def _log_metrics(self, trainer, step, metric_items):
         # pytorch-lightning runs a few steps of validation in the beginning of training
         # as a sanity check to catch bugs without having to wait for the training routine
         # to complete. During this check, we should skip logging metrics.
@@ -80,13 +94,24 @@ def _log_metrics(self, trainer, pl_module):
         if sanity_checking:
             return
 
+        # Cast metric value as  float before passing into logger.
+        metrics = dict(map(lambda x: (x[0], float(x[1])), metric_items))
+        self.metrics_logger.record_metrics(metrics, step)
+
+    def _log_epoch_metrics(self, trainer, pl_module):
+        # `trainer.callback_metrics` contains both training and validation metrics
+        # and includes metrics logged on steps and epochs.
+        # If we have logged any metrics on a step basis in mlflow, we exclude these from the
+        # epoch level metrics to prevent mixing epoch and step based values.
+        metric_items = [
+            (name, val)
+            for (name, val) in trainer.callback_metrics.items()
+            if name not in self._step_metrics
+        ]
+        # Record which metrics are logged on epochs, so we don't try to log these on steps
+        self._epoch_metrics.update(name for (name, _) in metric_items)
         if (pl_module.current_epoch + 1) % self.log_every_n_epoch == 0:
-            # `trainer.callback_metrics` contains both training and validation metrics
-            cur_metrics = trainer.callback_metrics
-            # Cast metric value as  float before passing into logger.
-            metrics = dict(map(lambda x: (x[0], float(x[1])), cur_metrics.items()))
-
-            self.metrics_logger.record_metrics(metrics, pl_module.current_epoch)
+            self._log_metrics(trainer, pl_module.current_epoch, metric_items)
 
     _pl_version = Version(pl.__version__)
 
@@ -100,7 +125,7 @@ def _log_metrics(self, trainer, pl_module):
         def on_train_epoch_end(
             self, trainer, pl_module, *args
         ):  # pylint: disable=signature-differs,arguments-differ,unused-argument
-            self._log_metrics(trainer, pl_module)
+            self._log_epoch_metrics(trainer, pl_module)
 
     # In pytorch-lightning >= 1.2.0, logging metrics in `on_epoch_end` results in duplicate
     # metrics records because `on_epoch_end` is called after both train and validation
@@ -128,7 +153,7 @@ def on_train_epoch_end(
             # log metrics in `on_validaion_epoch_end` to avoid logging the same metrics
             # records twice
             if not trainer.enable_validation:
-                self._log_metrics(trainer, pl_module)
+                self._log_epoch_metrics(trainer, pl_module)
 
         @rank_zero_only
         def on_validation_epoch_end(self, trainer, pl_module):
@@ -138,7 +163,7 @@ def on_validation_epoch_end(self, trainer, pl_module):
             :param trainer: pytorch lightning trainer instance
             :param pl_module: pytorch lightning base module
             """
-            self._log_metrics(trainer, pl_module)
+            self._log_epoch_metrics(trainer, pl_module)
 
     else:
 
@@ -150,7 +175,35 @@ def on_epoch_end(self, trainer, pl_module):
             :param trainer: pytorch lightning trainer instance
             :param pl_module: pytorch lightning base module
             """
-            self._log_metrics(trainer, pl_module)
+            self._log_epoch_metrics(trainer, pl_module)
+
+    @rank_zero_only
+    def on_train_batch_end(
+        self, trainer, pl_module, *args
+    ):  # pylint: disable=signature-differs,arguments-differ,unused-argument
+        """
+        Log metric values after each step
+
+        :param trainer: pytorch lightning trainer instance
+        :param pl_module: pytorch lightning base module
+        """
+        if not self.log_every_n_step:
+            return
+        # When logging at the end of a batch step, we only want to log metrics that are logged
+        # on steps. For forked metrics (metrics logged on both steps and epochs), we exclude the
+        # metric with the non-forked name (eg. "loss" when we have "loss", "loss_step" and
+        # "loss_epoch") so that this is only logged on epochs. We also record which metrics
+        # we've logged per step, so we can later exclude these from metrics logged on epochs.
+        metrics = _get_step_metrics(trainer)
+        metric_items = [
+            (name, val)
+            for (name, val) in metrics.items()
+            if (name not in self._epoch_metrics) and (f"{name}_step" not in metrics.keys())
+        ]
+        self._step_metrics.update(name for (name, _) in metric_items)
+        step = trainer.global_step
+        if (step + 1) % self.log_every_n_step == 0:
+            self._log_metrics(trainer, step, metric_items)
 
     @rank_zero_only
     def on_train_start(self, trainer, pl_module):
@@ -211,6 +264,22 @@ def on_test_end(self, trainer, pl_module):
         self.metrics_logger.flush()
 
 
+# PyTorch-Lightning refactored the LoggerConnector class in version 1.4.0 and made metrics
+# update on demand. Prior to this, the metrics from the current step were not available to
+# callbacks immediately, so the view of metrics was off by one step.
+# To avoid this problem, we access the metrics via the logger_connector for older versions.
+if _pl_version >= Version("1.4.0"):
+
+    def _get_step_metrics(trainer):
+        return trainer.callback_metrics
+
+
+else:
+
+    def _get_step_metrics(trainer):
+        return trainer.logger_connector.cached_results.get_latest_batch_log_metrics()
+
+
 def _log_early_stop_params(early_stop_callback, client, run_id):
     """
     Logs early stopping configuration parameters to MLflow.
@@ -276,6 +345,7 @@ def patched_fit(original, self, *args, **kwargs):
 
     log_models = get_autologging_config(mlflow.pytorch.FLAVOR_NAME, "log_models", True)
     log_every_n_epoch = get_autologging_config(mlflow.pytorch.FLAVOR_NAME, "log_every_n_epoch", 1)
+    log_every_n_step = get_autologging_config(mlflow.pytorch.FLAVOR_NAME, "log_every_n_step", None)
 
     early_stop_callback = None
     for callback in self.callbacks:
@@ -285,7 +355,9 @@ def patched_fit(original, self, *args, **kwargs):
 
     if not any(isinstance(callbacks, __MLflowPLCallback) for callbacks in self.callbacks):
         self.callbacks += [
-            __MLflowPLCallback(client, metrics_logger, run_id, log_models, log_every_n_epoch)
+            __MLflowPLCallback(
+                client, metrics_logger, run_id, log_models, log_every_n_epoch, log_every_n_step
+            )
         ]
 
     client.flush(synchronous=False)


Test patch:
diff --git a/tests/pytorch/iris.py b/tests/pytorch/iris.py
--- a/tests/pytorch/iris.py
+++ b/tests/pytorch/iris.py
@@ -48,6 +48,7 @@ def training_step(self, batch, batch_idx):
         self.train_acc(torch.argmax(logits, dim=1), y)
         self.log("train_acc", self.train_acc.compute(), on_step=False, on_epoch=True)
         self.log("loss", loss)
+        self.log("loss_forked", loss, on_epoch=True, on_step=True)
         return {"loss": loss}
 
     def validation_step(self, batch, batch_idx):
diff --git a/tests/pytorch/test_pytorch_autolog.py b/tests/pytorch/test_pytorch_autolog.py
--- a/tests/pytorch/test_pytorch_autolog.py
+++ b/tests/pytorch/test_pytorch_autolog.py
@@ -9,6 +9,7 @@
 from pytorch_lightning.callbacks import ModelCheckpoint
 from mlflow.utils.file_utils import TempDir
 from iris_data_module import IrisDataModule, IrisDataModuleWithoutValidation
+from mlflow.exceptions import MlflowException
 from mlflow.pytorch._pytorch_autolog import _get_optimizer_name
 from mlflow.tracking.client import MlflowClient
 
@@ -43,6 +44,20 @@ def pytorch_model_without_validation():
     return trainer, run
 
 
+@pytest.fixture(params=[(1, 1), (1, 10), (2, 1)])
+def pytorch_model_with_steps_logged(request):
+    log_every_n_epoch, log_every_n_step = request.param
+    mlflow.pytorch.autolog(log_every_n_epoch=log_every_n_epoch, log_every_n_step=log_every_n_step)
+    model = IrisClassification()
+    dm = IrisDataModule()
+    dm.setup(stage="fit")
+    trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
+    trainer.fit(model, dm)
+    client = mlflow.tracking.MlflowClient()
+    run = client.get_run(client.list_run_infos(experiment_id="0")[0].run_id)
+    return trainer, run, log_every_n_epoch, log_every_n_step
+
+
 @pytest.mark.parametrize("log_models", [True, False])
 def test_pytorch_autolog_log_models_configuration(log_models):
     mlflow.pytorch.autolog(log_models=log_models)
@@ -73,9 +88,19 @@ def test_pytorch_autolog_logs_expected_data(pytorch_model):
     _, run = pytorch_model
     data = run.data
 
-    # Checking if metrics are logged
+    # Checking if metrics are logged.
+    # When autolog is configured with the default configuration to not log on steps,
+    # then all metrics are logged per epoch, including step based metrics.
     client = mlflow.tracking.MlflowClient()
-    for metric_key in ["loss", "train_acc", "val_loss", "val_acc"]:
+    for metric_key in [
+        "loss",
+        "train_acc",
+        "val_loss",
+        "val_acc",
+        "loss_forked",
+        "loss_forked_step",
+        "loss_forked_epoch",
+    ]:
         assert metric_key in run.data.metrics
         metric_history = client.get_metric_history(run.info.run_id, metric_key)
         assert len(metric_history) == NUM_EPOCHS
@@ -102,6 +127,50 @@ def test_pytorch_autolog_logs_expected_metrics_without_validation(pytorch_model_
         assert len(metric_history) == NUM_EPOCHS
 
 
+@pytest.mark.skipif(
+    Version(pl.__version__) < Version("1.1.0"),
+    reason="Access to metrics from the current step is only possible since PyTorch-lightning 1.1.0"
+    "when LoggerConnector.cached_results was added",
+)
+def test_pytorch_autolog_logging_forked_metrics_on_step_and_epoch(
+    pytorch_model_with_steps_logged,
+):
+    # When autolog is configured to log on steps as well as epochs,
+    # then we only log step based metrics per step and not on epochs.
+    trainer, run, log_every_n_epoch, log_every_n_step = pytorch_model_with_steps_logged
+    num_logged_steps = trainer.global_step // log_every_n_step
+    num_logged_epochs = NUM_EPOCHS // log_every_n_epoch
+
+    client = mlflow.tracking.MlflowClient()
+    for (metric_key, expected_len) in [
+        ("train_acc", num_logged_epochs),
+        ("loss", num_logged_steps),
+        ("loss_forked", num_logged_epochs),
+        ("loss_forked_step", num_logged_steps),
+        ("loss_forked_epoch", num_logged_epochs),
+    ]:
+        assert metric_key in run.data.metrics, f"Missing {metric_key} in metrics"
+        metric_history = client.get_metric_history(run.info.run_id, metric_key)
+        assert (
+            len(metric_history) == expected_len
+        ), f"Expected {expected_len} values for {metric_key}, got {len(metric_history)}"
+
+
+@pytest.mark.skipif(
+    Version(pl.__version__) >= Version("1.1.0"),
+    reason="Logging step metrics is supported since PyTorch-Lightning 1.1.0",
+)
+def test_pytorch_autolog_raises_error_when_step_logging_unsupported():
+    mlflow.pytorch.autolog(log_every_n_step=1)
+    model = IrisClassification()
+    dm = IrisDataModule()
+    trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
+    with pytest.raises(
+        MlflowException, match="log_every_n_step is only supported for PyTorch-Lightning >= 1.1.0"
+    ):
+        trainer.fit(model, dm)
+
+
 # pylint: disable=unused-argument
 def test_pytorch_autolog_persists_manually_created_run():
     with mlflow.start_run() as manual_run: