Repository: Lightning-AI/pytorch-lightning
Pull request: feat(wandb): offset logging step when resuming
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
-->

When resuming a run, the step was not being offset to account for previous logged data.
It will now be correctly offset.

In addition, I took advantage to give users a hint when they have related issues trying to log "back in time". It is usually due to them logging manually other data (images, etc) and causing the wandb step to automatically increase. This will close https://github.com/wandb/client/issues/1507

I added the same details in the documentation.

## Before submitting
- [X] Was this discussed/approved via a Github issue? discussed in wandb issue
- [X] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [X] Did you make sure your PR does only one thing, instead of bundling different changes together? Since the other modification is just a warning
- [X] Did you make sure to update the documentation with your changes?
- [X] Did you write any new necessary tests? 
- [X] Did you verify new and existing tests pass locally with your changes?
- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:  

 - [ ] Is this pull request ready for review? (if not, please submit in draft mode)
 - [ ] Check that all items from **Before submitting** are resolved
 - [ ] Make sure the title is self-explanatory and the description concisely explains the PR
 - [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified; _Bugfixes should be including in bug-fix release milestones (m.f.X) and features should be included in (m.X.b) releases._
 

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
WandB-Logger drops all the logged values in training step for PyTorch Lightning
train-step:
```
def training_step(self, batch, batch_idx):
            x, y = batch
            logits, masks = self(x)
            loss = self.loss_func(logits, y)
            self.log(f"train-loss", loss, prog_bar=True)
            return {
                "loss": loss,
                "pred": torch.argmax(logits, -1),
                "labels": y,
            }
```
and on epoch end it logs the Precision and recall. 
after the validation step. I get a ` wandb: WARNING Step must only increase in log calls.  Step 379 < 380; dropping {train-loss: 0.494, 'train-precision': 0.9, 'train-recall': 1, 'epoch': 1}`

WandB version: 0.10.10
PyTorch-Lightning: 1.0.6

Hints from issue discussion:
Hi @IsCoelacanth 
How are you logging the other parameters (precision and recall)?
I have this issue too! PyTorch Lightning resets step to 0 at the start of a new epoch, but wandb expects step to keep increasing and seems to ignore that we're on a new epoch. I don't know how to resolve this, it seems like a mismatch of expectations for the behavior of step.
@isaacrob step works as expected with me (increasing even when changing epoch). See [this colab](https://colab.research.google.com/drive/16d1uctGaw2y9KhGBlINNTsWpmlXdJwRW)
Interesting. Another issue I'm facing is that epochs is being logged as a chart, with per-step reporting and everything, which doesn't make any sense and doesn't seem to be happening in your notebook based off the screenshot. Here's a screenshot of mine: 
![Screenshot from 2020-11-19 17-17-30](https://user-images.githubusercontent.com/6684952/99743817-4e745680-2a8b-11eb-9ffa-73edb5c2ec43.png)

I wonder if this is related to it not handling steps during change of epoch correctly?

Yes epoch is now a constant (it comes directly from Pytorch Lightning) though it should increase over time.
It looks like you only have one epoch, unless you try to reset the logger at each epoch.

See what I get from the colab I shared: [W&B run](https://wandb.ai/borisd13/MNIST/runs/1tjej3u4?workspace=user-borisd13)
Ok so that's not it. Like OP I'm getting this warning at the end of my first validation batch:
`wandb: WARNING Step must only increase in log calls.  Step 199 < 201; dropping {'val_cls_loss': 1.3014671802520752, 'val_box_loss': 0.9644062519073486, 'val_masking_loss': 0.544899582862854, 'val_tot_loss': 2.8107728958129883, 'epoch': 0}.`
and I don't see any metrics related to validation show up online, which appears to be because it's dropping all of them. It looks like then that step might be reset during validation call? When I reach the next epoch the same thing happens:
`wandb: WARNING Step must only increase in log calls.  Step 399 < 401; dropping {'val_cls_loss': 1.0469279289245605, 'val_box_loss': 0.9169614315032959, 'val_masking_loss': 0.49405547976493835, 'val_tot_loss': 2.457944869995117, 'epoch': 1}.`
Fixed it! I was calling `trainer.logger.experiment.log` in a callback to log Matplotlib images and apparently this increases the step unless `commit` is set to `False`, in which case it gets group with the next logging message sent via PyTorch Lightning. Subtle bug! Everything's ok now and working as expected :) thanks for your help @borisdayma. Hope @IsCoelacanth was able to resolve their issue too!
hey, @borisdayma there's a callback on validation completion which logs the metrics calculated using SKLearn metrics. same way @isaacrob mentioned. 
setting `commit = False` removes the warning and behaves the same way as mentioned by Isaac. 
When you set `commit=False` or when you explicitly set a `step`, the values are aggregated and sent only when step increments (or at the end of the run) so you don't see them immediately in the dashboard.

Here is the logic:
* `wandb.log(dict, step=50)`
called automatically by `WandLogger` during training and validation
* **Option 1**: manual call to `wandb.log(my_dict)`
this will log at step 50, and increment to step 51 ; if you do multiple manual logs, step will increase, causing an issue if you try to log "back in time"
* **Option 2**: manual call to `wandb.log(my_dict, commit=False)`
this will log at step 50 and not increment
* **Option 3**: manual call to `wandb.log(my_dict, step=trainer.global_step)`
this will ensure your logging is associated to the right step

In practice, just setting `commit=False` should work well since you will probably do manual logging after a training step, which means all the metrics will be aggregated to last training step logged (which I guess is what you want).

Let me know what you think.
We could output a message suggesting to use `commit=False` on manual logging when this situation occurs.

Commit message: feat: improve warning

Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>

Patch to review:
diff --git a/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py
index 24007c3a04307..7b3cf26403717 100644
--- a/pytorch_lightning/loggers/wandb.py
+++ b/pytorch_lightning/loggers/wandb.py
@@ -31,6 +31,7 @@
 
 from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
 from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning import _logger as log
 
 
 class WandbLogger(LightningLoggerBase):
@@ -66,6 +67,9 @@ class WandbLogger(LightningLoggerBase):
         wandb_logger = WandbLogger()
         trainer = Trainer(logger=wandb_logger)
 
+    Note: When logging manually through `wandb.log` or `trainer.logger.experiment.log`,
+    make sure to use `commit=False` so the logging step does not increase.
+
     See Also:
         - `Tutorial <https://app.wandb.ai/cayush/pytorchlightning/reports/
           Use-Pytorch-Lightning-with-Weights-%26-Biases--Vmlldzo2NjQ1Mw>`__
@@ -74,6 +78,7 @@ class WandbLogger(LightningLoggerBase):
     """
 
     LOGGER_JOIN_CHAR = '-'
+    WARNING_DISPLAYED = False
 
     def __init__(
         self,
@@ -103,7 +108,7 @@ def __init__(
         self._log_model = log_model
         self._prefix = prefix
         self._kwargs = kwargs
-        # logging multiple Trainer on a single W&B run (k-fold, etc)
+        # logging multiple Trainer on a single W&B run (k-fold, resuming, etc)
         self._step_offset = 0
 
     def __getstate__(self):
@@ -134,6 +139,8 @@ def experiment(self) -> Run:
             self._experiment = wandb.init(
                 name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous,
                 id=self._id, resume='allow', **self._kwargs) if wandb.run is None else wandb.run
+            # offset logging step when resuming a run
+            self._step_offset = self._experiment.step
             # save checkpoints in wandb dir to upload on W&B servers
             if self._log_model:
                 self._save_dir = self._experiment.dir
@@ -154,6 +161,9 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
         assert rank_zero_only.rank == 0, 'experiment tried to log from global_rank != 0'
 
         metrics = self._add_prefix(metrics)
+        if not WandbLogger.WARNING_DISPLAYED and (step is not None and step + self._step_offset < self.experiment.step):
+            log.warning('Trying to log at a previous step. Use `commit=False` when logging metrics manually.')
+            WandbLogger.WARNING_DISPLAYED = True
         self.experiment.log(metrics, step=(step + self._step_offset) if step is not None else None)
 
     @property
diff --git a/tests/loggers/test_all.py b/tests/loggers/test_all.py
index ba5791c7b9f4a..89c731d432ee9 100644
--- a/tests/loggers/test_all.py
+++ b/tests/loggers/test_all.py
@@ -74,7 +74,9 @@ def test_loggers_fit_test_all(tmpdir, monkeypatch):
     with mock.patch('pytorch_lightning.loggers.test_tube.Experiment'):
         _test_loggers_fit_test(tmpdir, TestTubeLogger)
 
-    with mock.patch('pytorch_lightning.loggers.wandb.wandb'):
+    with mock.patch('pytorch_lightning.loggers.wandb.wandb') as wandb:
+        wandb.run = None
+        wandb.init().step = 0
         _test_loggers_fit_test(tmpdir, WandbLogger)
 
 
@@ -368,5 +370,7 @@ def test_logger_with_prefix_all(tmpdir, monkeypatch):
     # WandB
     with mock.patch('pytorch_lightning.loggers.wandb.wandb') as wandb:
         logger = _instantiate_logger(WandbLogger, save_idr=tmpdir, prefix=prefix)
+        wandb.run = None
+        wandb.init().step = 0
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.log.assert_called_once_with({'tmp-test': 1.0}, step=0)
diff --git a/tests/loggers/test_wandb.py b/tests/loggers/test_wandb.py
index 33211e6492d91..708f15cfb2734 100644
--- a/tests/loggers/test_wandb.py
+++ b/tests/loggers/test_wandb.py
@@ -34,6 +34,9 @@ def test_wandb_logger_init(wandb):
     wandb.init.assert_called_once()
     wandb.init().log.assert_called_once_with({'acc': 1.0}, step=None)
 
+    # mock wandb step
+    wandb.init().step = 0
+
     # test wandb.init not called if there is a W&B run
     wandb.init().log.reset_mock()
     wandb.init.reset_mock()
@@ -71,6 +74,7 @@ def test_wandb_pickle(wandb, tmpdir):
     class Experiment:
         """ """
         id = 'the_id'
+        step = 0
 
         def project_name(self):
             return 'the_project_name'
@@ -108,8 +112,11 @@ def test_wandb_logger_dirs_creation(wandb, tmpdir):
     assert logger.name is None
 
     # mock return values of experiment
+    wandb.run = None
+    wandb.init().step = 0
     logger.experiment.id = '1'
     logger.experiment.project_name.return_value = 'project'
+    logger.experiment.step = 0
 
     for _ in range(2):
         _ = logger.experiment


Code review comments:
diff_hunk: @@ -154,6 +161,9 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
         assert rank_zero_only.rank == 0, 'experiment tried to log from global_rank != 0'
 
         metrics = self._add_prefix(metrics)
+        if not WandbLogger.WARNING_DISPLAYED and (step is not None and step + self._step_offset < self.experiment.step):
+            log.warning('Trying to log at a previous step. Use `commit=False` when logging metrics manually.')
+            WandbLogger.WARNING_DISPLAYED = True
path: pytorch_lightning/loggers/wandb.py
text: @user1:
```suggestion
        if not self.WARNING_DISPLAYED and (step is not None and step + self._step_offset < self.experiment.step):
            log.warning('Trying to log at a previous step. Use `commit=False` when logging metrics manually.')
            self.WARNING_DISPLAYED = True
```