Lightning-AI / Lightning-AI/pytorch-lightning

Trainer's current_epoch is not updated at on_fit_start when loading from a checkpoint

Open
#17,712 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

### Bug description

When reloading from a checkpoint (using `Trainer().fit(..., ckpt_path=X)`) it seems that the `current_epoch` of the trainer is not updated accordingly at `on_fit_start`, but rather later on at `on_train_epoch_end`.

`load_state_dict` is called already for callbacks, so we know it is indeed a training that is loaded from a ckpt, not from scratch.

My workaround for this (I have some custom logic in a callback that only applies if we continue/fine-tune a training) is to look at `trainer.ckpt_path` not being `None`.

Is this the intended behavior or should `trainer.current_epoch` be also >0 ?

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

v2.0

### How to reproduce the bug

```python
from pytorch_lightning import LightningModule, Callback, Trainer
from pytorch_lightning.callbacks.callback import Callback
from torch.utils.data import Dataset, DataLoader
import torch as tr
from torch import nn
from torch.nn import functional as F

class Reader(Dataset):
def __getitem__(self, ix):
return tr.randn(5), tr.randn(1)

def __len__(self):
return 5

class MyCallback(Callback):
def on_train_epoch_start(self, trainer, pl_module):
print(f"[on_train_epoch_start] epoch: {trainer.current_epoch}")

def on_fit_start(self, trainer, pl_module):
print(f"[on_fit_start] epoch: {trainer.current_epoch}")

def load_state_dict(self, state):
print("[load_state_dict]")

def state_dict(self):
print("[state_dict]")
return {1: 2}

class LitModel(LightningModule):
def __init__(self):
super().__init__()
self.l1 = nn.Linear(5, 1)

def forward(self, x):
return tr.relu(self.l1(x.view(x.size(0), -1)))

def training_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
loss = F.cross_entropy(y_hat, y)
return loss

def configure_optimizers(self):
return tr.optim.Adam(self.parameters(), lr=0.02)

def main():
loader = DataLoader(Reader(), batch_size=5)
model = LitModel()
Trainer(max_epochs=2, callbacks=[MyCallback()], enable_progress_bar=False, enable_model_summary=False) \
.fit(model, train_dataloaders=loader)
print("----")
Trainer(max_epochs=5, callbacks=[MyCallback()], enable_progress_bar=False, enable_model_summary=False) \
.fit(model, train_dataloaders=loader, ckpt_path=model.trainer.checkpoint_callback.best_model_path)

if __name__ == "__main__":
main()
```

### Error messages and logs

```bash
python main.py 2>/dev/null

[on_fit_start] epoch: 0
[on_train_epoch_start] epoch: 0
[state_dict]
[on_train_epoch_start] epoch: 1
[state_dict]
----
[load_state_dict] <- we know that this is a loaded train since this method is called here
[on_fit_start] epoch: 0 <- here epoch is still 0, but later on it restarts from 2
[on_train_epoch_start] epoch: 2
[state_dict]
[on_train_epoch_start] epoch: 3
[state_dict]
[on_train_epoch_start] epoch: 4
[state_dict]

```

### Environment

Current environment

* CUDA:
- GPU: None
- available: False
- version: 11.7
* Lightning:
- lightning-module-enhanced: 0.1
- lightning-utilities: 0.8.0
- pytorch-lightning: 2.0.0
- torch: 2.0.0
- torchinfo: 1.7.0
- torchmetrics: 0.11.4
- torchvision: 0.15.1
* Packages:
- absl-py: 1.2.0
- aiohttp: 3.8.1
- aiosignal: 1.2.0
- alembic: 1.10.3
- antlr4-python3-runtime: 4.9.3
- anyio: 3.6.2
- appdirs: 1.4.4
- argon2-cffi: 21.3.0
- argon2-cffi-bindings: 21.2.0
- arrow: 1.2.3
- art: 5.8
- astroid: 2.15.5
- asttokens: 2.2.1
- astunparse: 1.6.3
- async-generator: 1.10
- async-timeout: 4.0.2
- attrs: 21.4.0
- audioread: 2.1.9
- automat: 22.10.0
- backcall: 0.2.0
- beautifulsoup4: 4.11.1
- bleach: 6.0.0
- cachetools: 5.2.0
- certifi: 2022.12.7
- cffi: 1.15.1
- charset-normalizer: 2.1.0
- cloudpickle: 2.2.1
- cmake: 3.26.3
- colorama: 0.4.6
- coloredlogs: 15.0.1
- comm: 0.1.2
- constantly: 15.1.0
- cryptography: 39.0.2
- cssselect: 1.2.0
- cycler: 0.11.0
- databricks-cli: 0.17.6
- debugpy: 1.6.6
- decorator: 5.1.1
- decord: 0.6.0
- defusedxml: 0.7.1
- dill: 0.3.5.1
- docker: 6.0.1
- entrypoints: 0.4
- exceptiongroup: 1.1.0
- executing: 1.2.0
- fastcore: 1.5.28
- fastjsonschema: 2.16.3
- ffmpeg-python: 0.2.0
- filelock: 3.7.1
- flask: 2.2.3
- flatbuffers: 1.12
- flow-vis: 0.1
- fonttools: 4.34.4
- fqdn: 1.5.1
- frozenlist: 1.3.0
- fsspec: 2022.5.0
- gast: 0.4.0
- gdown: 4.6.0
- gitdb: 4.0.10
- gitpython: 3.1.31
- google-auth: 2.9.1
- google-auth-oauthlib: 0.4.6
- google-pasta: 0.2.0
- graphviz: 0.20.1
- grpcio: 1.47.0
- gunicorn: 20.1.0
- h11: 0.14.0
- h5py: 3.7.0
- huggingface-hub: 0.12.0
- humanfriendly: 10.0
- hyperlink: 21.0.0
- idna: 3.4
- imageio: 2.19.5
- imageio-ffmpeg: 0.4.7
- incremental: 22.10.0
- iniconfig: 1.1.1
- ipykernel: 6.21.2
- ipython: 8.10.0
- ipython-genutils: 0.2.0
- ipywidgets: 8.0.4
- isoduration: 20.11.0
- isort: 5.10.1
- itemadapter: 0.7.0
- itemloaders: 1.0.6
- itsdangerous: 2.1.2
- jedi: 0.18.2
- jinja2: 3.1.2
- jmespath: 1.0.1
- joblib: 1.1.0
- jsonpointer: 2.3
- jsonschema: 4.17.3
- jupyter: 1.0.0
- jupyter-client: 8.0.3
- jupyter-console: 6.6.1
- jupyter-core: 5.2.0
- jupyter-events: 0.6.3
- jupyter-server: 2.3.0
- jupyter-server-terminals: 0.4.4
- jupyterlab-pygments: 0.2.2
- jupyterlab-widgets: 3.0.5
- keras: 2.9.0
- keras-preprocessing: 1.1.2
- kiwisolver: 1.4.4
- lazy-loader: 0.2
- lazy-object-proxy: 1.7.1
- libclang: 14.0.6
- librosa: 0.9.2
- lightning-module-enhanced: 0.1
- lightning-utilities: 0.8.0
- lit: 16.0.5
- llvmlite: 0.38.1
- lovely-numpy: 0.2.8
- lovely-tensors: 0.1.14
- lxml: 4.9.2
- lycon: 0.2.0
- markdown: 3.4.1
- markupsafe: 2.1.2
- matplotlib: 3.5.2
- matplotlib-inline: 0.1.6
- mccabe: 0.7.0
- mistune: 2.0.5
- mlflow: 2.2.2
- mpmath: 1.2.1
- multidict: 6.0.2
- natsort: 8.2.0
- nbclassic: 0.5.2
- nbclient: 0.7.2
- nbconvert: 7.2.9
- nbformat: 5.7.3
- nest-asyncio: 1.5.6
- networkx: 2.8.5
- notebook: 6.5.2
- notebook-shim: 0.2.2
- numba: 0.55.2
- numpy: 1.24.2
- nvidia-cublas-cu11: 11.10.3.66
- nvidia-cuda-cupti-cu11: 11.7.101
- nvidia-cuda-nvrtc-cu11: 11.7.99
- nvidia-cuda-runtime-cu11: 11.7.99
- nvidia-cudnn-cu11: 8.5.0.96
- nvidia-cufft-cu11: 10.9.0.58
- nvidia-curand-cu11: 10.2.10.91
- nvidia-cusolver-cu11: 11.4.0.1
- nvidia-cusparse-cu11: 11.7.4.91
- nvidia-nccl-cu11: 2.14.3
- nvidia-nvtx-cu11: 11.7.91
- nwgraph: 0.2.1
- nwutils: 0.3.4
- omegaconf: 2.3.0
- onnx: 1.13.1
- onnxruntime: 1.14.0
- opencv-python: 4.6.0.66
- opt-einsum: 3.3.0
- outcome: 1.2.0
- overrides: 7.3.1
- packaging: 21.3
- pandarallel: 1.6.4
- pandas: 1.5.1
- pandocfilters: 1.5.0
- parsel: 1.7.0
- parso: 0.8.3
- pexpect: 4.8.0
- pickleshare: 0.7.5
- pillow: 9.4.0
- pims: 0.6.1
- pip: 23.1.2
- platformdirs: 2.5.2
- pluggy: 1.0.0
- plyvel: 1.5.0
- pooch: 1.6.0
- pool-resources: 0.1
- prometheus-client: 0.16.0
- prompt-toolkit: 3.0.37
- protego: 0.2.1
- protobuf: 3.20.1
- psutil: 5.9.5
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- py: 1.11.0
- pyarrow: 11.0.0
- pyasn1: 0.4.8
- pyasn1-modules: 0.2.8
- pycm: 3.8
- pycparser: 2.21
- pydeprecate: 0.3.2
- pydispatcher: 2.0.7
- pygments: 2.14.0
- pylint: 2.17.4
- pyopenssl: 23.0.0
- pyparsing: 3.0.9
- pyproj: 3.4.1
- pyrsistent: 0.19.3
- pysocks: 1.7.1
- pytest: 7.3.1
- python-dateutil: 2.8.2
- python-json-logger: 2.0.7
- pytorch-lightning: 2.0.0
- pytz: 2023.3
- pywavelets: 1.3.0
- pyyaml: 6.0
- pyzmq: 25.0.0
- qtconsole: 5.4.0
- qtpy: 2.3.0
- querystring-parser: 1.2.4
- queuelib: 1.6.2
- requests: 2.28.2
- requests-file: 1.5.1
- requests-oauthlib: 1.3.1
- resampy: 0.3.1
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rsa: 4.9
- scikit-image: 0.20.0
- scikit-learn: 1.1.1
- scipy: 1.8.1
- scrapy: 2.8.0
- selenium: 4.8.2
- send2trash: 1.8.0
- service-identity: 21.1.0
- setuptools: 65.6.3
- shap: 0.41.0
- six: 1.16.0
- slicer: 0.0.7
- slicerator: 1.1.0
- smmap: 5.0.0
- sniffio: 1.3.0
- sortedcontainers: 2.4.0
- soundfile: 0.10.3.post1
- soupsieve: 2.3.2.post1
- sqlalchemy: 2.0.9
- sqlparse: 0.4.3
- stack-data: 0.6.2
- sympy: 1.11.1
- tabulate: 0.9.0
- tensorboard-data-server: 0.6.1
- tensorboard-plugin-wit: 1.8.1
- tensorboardx: 2.5.1
- tensorflow: 2.9.0
- tensorflow-estimator: 2.9.0
- tensorflow-io-gcs-filesystem: 0.29.0
- termcolor: 2.1.1
- terminado: 0.17.1
- tf-estimator-nightly: 2.8.0.dev2021122109
- threadpoolctl: 3.1.0
- tifffile: 2022.5.4
- timm: 0.6.12
- tinycss2: 1.2.1
- tldextract: 3.4.0
- tomli: 2.0.1
- tomlkit: 0.11.1
- torch: 2.0.0
- torchinfo: 1.7.0
- torchmetrics: 0.11.4
- torchvision: 0.15.1
- tornado: 6.2
- tqdm: 4.65.0
- traitlets: 5.9.0
- transforms3d: 0.4.1
- trio: 0.22.0
- trio-websocket: 0.9.2
- triton: 2.0.0
- twisted: 22.10.0
- typing-extensions: 4.3.0
- typing-utils: 0.1.0
- uri-template: 1.2.0
- urllib3: 1.26.15
- w3lib: 2.1.1
- wcwidth: 0.2.6
- webcolors: 1.12
- webencodings: 0.5.1
- websocket-client: 1.5.1
- werkzeug: 2.2.3
- wheel: 0.37.1
- widgetsnbextension: 4.0.5
- wrapt: 1.14.1
- wsproto: 1.2.0
- yarl: 1.7.2
- yt-dlp: 2023.3.4
- zope.interface: 5.5.2
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.11
- release: 5.19.0-41-generic
- version: #42~22.04.1-Ubuntu SMP PREEMPT_DYNAMIC Tue Apr 18 17:40:00 UTC 2

### More info

_No response_

cc @lantiga

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 with the Trainer.fit checkpoint-loading path and the callback lifecycle around load_state_dict and on_fit_start, using the supplied reproduction script. Trace where current_epoch is restored, then add coverage showing that a resumed fit exposes the checkpoint epoch during on_fit_start and still resumes at the correct training epoch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.