Different data-loading behaviour in 1.6.0 version.

🐛 Bug

I observe different data-loading behaviour after upgrading PLT to 1.6.0 or 1.6.1 version (1.5.10 works fine). I'm running nn-UNet from DeepLearningExamples repository. Model is using NVIDIA DALI data loader.

Problem description:

The first epoch is running fine, but the second epoch is counted as third instead of second. Also it looks like the validation runs over empty set (starting from the second epoch), as I'm getting warning /opt/conda/lib/python3.8/site-packages/torchmetrics/utilities/prints.py:36: UserWarning: The ``compute`` method of metric Dice was called before the update method which may lead to errors, as metric states have not yet been updated. warnings.warn(*args, **kwargs) after the first epoch finishes.

I wonder if this can be because of Re-define the current_epoch boundary or Re-define the global_step boundary updates in 1.6.0.

The log in 1.6.0 I see is:

DLL 2022-04-20 10:20:03.907821 - 0 Mean dice : 78.96  Highest : 78.96  L1 : 79.04  L2 : 78.88  val_loss : 0.3151 
Metric dice_mean improved. New best score: 78.961
/opt/conda/lib/python3.8/site-packages/torchmetrics/utilities/prints.py:36: UserWarning: The ``compute`` method of metric Dice was called before the update method which may lead to errors, as metric states have not yet been updated.
  warnings.warn(*args, **kwargs)
DLL 2022-04-20 10:20:09.509141 - 2 Mean dice : nan  Highest : 78.96  L1 : nan  L2 : nan  val_loss : nan 
Monitored metric dice_mean = nan is not finite. Previous best value was 78.961. Signaling Trainer to stop.
Trainer was signaled to stop but required minimum epochs (1000) or minimum steps (None) has not been met. Training will continue...
Monitored metric dice_mean = nan is not finite. Previous best value was 78.961. Signaling Trainer to stop.
Trainer was signaled to stop but required minimum epochs (1000) or minimum steps (None) has not been met. Training will continue...
DLL 2022-04-20 10:20:15.622250 - 4 Mean dice : 84.51  Highest : 84.51  L1 : 84.63  L2 : 84.38  val_loss : 0.164 
Metric dice_mean improved by 5.545 >= min_delta = 0.0. New best score: 84.507
DLL 2022-04-20 10:20:21.219675 - 6 Mean dice : nan  Highest : 84.51  L1 : nan  L2 : nan  val_loss : nan 
Monitored metric dice_mean = nan is not finite. Previous best value was 84.507. Signaling Trainer to stop.
Trainer was signaled to stop but required minimum epochs (1000) or minimum steps (None) has not been met. Training will continue...
Monitored metric dice_mean = nan is not finite. Previous best value was 84.507. Signaling Trainer to stop.
Trainer was signaled to stop but required minimum epochs (1000) or minimum steps (None) has not been met. Training will continue...

And the expected log from 1.5.10 PLT version is:

DLL 2022-04-20 10:18:36.048561 - 0 Mean dice : 78.56  Highest : 78.56  L1 : 78.61  L2 : 78.51  val_loss : 0.3184 
Metric dice_mean improved. New best score: 78.563
DLL 2022-04-20 10:18:42.251821 - 1 Mean dice : 80.61  Highest : 80.61  L1 : 79.66  L2 : 81.57  val_loss : 0.2183 
Metric dice_mean improved by 2.052 >= min_delta = 0.0. New best score: 80.615
DLL 2022-04-20 10:18:48.447890 - 2 Mean dice : 84.67  Highest : 84.67  L1 : 85.2  L2 : 84.13  val_loss : 0.1626 
Metric dice_mean improved by 4.054 >= min_delta = 0.0. New best score: 84.669
DLL 2022-04-20 10:18:54.646806 - 3 Mean dice : 83.34  Highest : 84.67  L1 : 82.5  L2 : 84.19  val_loss : 0.1689 
DLL 2022-04-20 10:19:00.841899 - 4 Mean dice : 82.37  Highest : 84.67  L1 : 83.09  L2 : 81.66  val_loss : 0.1729 
DLL 2022-04-20 10:19:07.035390 - 5 Mean dice : 83.97  Highest : 84.67  L1 : 85.05  L2 : 82.89  val_loss : 0.1578 
DLL 2022-04-20 10:19:13.227074 - 6 Mean dice : 85.69  Highest : 85.69  L1 : 86.4  L2 : 84.98  val_loss : 0.1415 
Metric dice_mean improved by 1.021 >= min_delta = 0.0. New best score: 85.690

To Reproduce

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/Segmentation/nnUNet
docker build -t nnunet .
mkdir data results
docker run -it --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/data:/data -v ${PWD}/results:/results nnunet:latest /bin/bash
python download.py --task 04
python preprocess.py --task 04 --dim 3
python main.py --task 04 --dim 3 --gpus 1

Expected behavior

Epochs are counted without gaps and validation is running on full validation set.

Environment

• CUDA:
  - GPU:
  - NVIDIA TITAN V
  - available: True
  - version: 11.6
• Packages:
  - numpy: 1.22.3
  - pyTorch_debug: False
  - pyTorch_version: 1.12.0a0+bd13bc6
  - pytorch-lightning: 1.6.1
  - tqdm: 4.63.1
• System:
  - OS: Linux
  - architecture:
  - 64bit
  - ELF
  - processor: x86_64
  - python: 3.8.13
  - version: #119~18.04.1-Ubuntu SMP Tue Mar 8 11:21:24 UTC 2022

Reproduce the behavior in DeepLearningExamples/PyTorch/Segmentation/nnUNet using the listed Docker and training commands, comparing pytorch-lightning 1.5.10 with 1.6.1. Start at main.py and investigate the epoch and validation behavior around the reported current_epoch and global_step boundary changes. Done means epochs have no gaps and validation runs over the full validation set.