Lightning-AI / Lightning-AI/pytorch-lightning

Metrics get mapped twice to the same epoch in MLflow logger

Open
#20,902 3 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug logger logger: mlflow ver: 2.5.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

The MLFlow logger associates two consecutive values of a metric with one epoch in the metric plots in the MLFlow ui.
This does not happen when x-axis is `steps`.

I'm logging my metrics by using `on_step=False` and `on_epoch=True` using the following code:

```
def training_step(self, batch: torch.Tensor, batch_idx: int):
outdict = self.model_step(batch.flatten(start_dim=1).to(dtype=torch.float32))

# Decide what to log:
self.log_dict(
self._log_dict(outdict, "train", dataloader_idx=0),
prog_bar=False,
on_step=False,
on_epoch=True,
logger=True,
sync_dist=True,
add_dataloader_idx=False,
)
return outdict

def validation_step(
self, batch: torch.Tensor, batch_idx: int, dataloader_idx: Optional[int] = 0
):
outdict = self.model_step(batch.flatten(start_dim=1).to(dtype=torch.float32))
self.log_dict(
self._log_dict(outdict, "val", dataloader_idx=dataloader_idx),
prog_bar=False,
on_step=False,
on_epoch=True,
logger=True,
sync_dist=True,
add_dataloader_idx=False,
)

return outdict
```
I'm not calling `log_metric` or `log_dict` anywhere else.

The plots look like this when having `step` on the x-axis.

Image

The same plot using `epoch` on the x-axis.

Image

The epoch vs step plot looks fine though

Image

Ignore the last point, that is due to running on the test data. Even if I don't run on this data, I still get the 'double points' for each epoch.

Am I doing something wrong? This seems like a pretty basic use case.

### What version are you seeing the problem on?

v2.5

### Reproduced in studio

_No response_

### How to reproduce the bug

```python

```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

### More info

_No response_

cc @lantiga @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the behavior with the shown training_step and log_dict configuration, focusing on the MLflow logger and epoch-based metric mapping. Inspect the logger path handling on_epoch=True and compare it with steps-axis logging; done when each epoch produces one metric point without changing the correct steps-axis behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
observability
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.