Lightning-AI / Lightning-AI/pytorch-lightning

trainer.test(ckpt_path='best') does not work as expected

Open
#17,312 15 comments 6 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

### Bug description

`trainer.test(model=model, ckpt_path='best')` works after `trainer.fit` but not otherwise

We get ```ValueError: `.test(ckpt_path="best")` is set but `ModelCheckpoint` is not configured to save the best model.```

`ModelCheckpoint` is configured to save the best model. In fact `save_top_k=1` is the default
```
checkpoint_callback = ModelCheckpoint(
dirpath='/Users/adam.amster/Downloads/pl_test',
monitor='val_loss',
save_top_k=1
)
```

Therefore the best model was saved and `ckpt_path="best"` should always work, regardless of if it's called after `fit`

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

2.0+

### How to reproduce the bug

```python
import os

import lightning as L
import torch
from lightning.pytorch.callbacks import ModelCheckpoint
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchmetrics import Accuracy
from torchvision import transforms
from torchvision.datasets import MNIST

PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
BATCH_SIZE = 256 if torch.cuda.is_available() else 64

class LitMNIST(L.LightningModule):
def __init__(self, data_dir=PATH_DATASETS, hidden_size=64, learning_rate=2e-4):
super().__init__()

# Set our init args as class attributes
self.data_dir = data_dir
self.hidden_size = hidden_size
self.learning_rate = learning_rate

# Hardcode some dataset specific attributes
self.num_classes = 10
self.dims = (1, 28, 28)
channels, width, height = self.dims
self.transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
]
)

# Define PyTorch model
self.model = nn.Sequential(
nn.Flatten(),
nn.Linear(channels * width * height, hidden_size),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden_size, self.num_classes),
)

self.val_accuracy = Accuracy(task="multiclass", num_classes=10)
self.test_accuracy = Accuracy(task="multiclass", num_classes=10)

def forward(self, x):
x = self.model(x)
return F.log_softmax(x, dim=1)

def training_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)

self.log("train_loss", loss)

return loss

def validation_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)
preds = torch.argmax(logits, dim=1)
self.val_accuracy.update(preds, y)

# Calling self.log will surface up scalars for you in TensorBoard
self.log("val_loss", loss, prog_bar=True)
self.log("val_acc", self.val_accuracy, prog_bar=True)

def test_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)
preds = torch.argmax(logits, dim=1)
self.test_accuracy.update(preds, y)

# Calling self.log will surface up scalars for you in TensorBoard
self.log("test_loss", loss, prog_bar=True)
self.log("test_acc", self.test_accuracy, prog_bar=True)

def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
return optimizer

####################
# DATA RELATED HOOKS
####################

def prepare_data(self):
# download
MNIST(self.data_dir, train=True, download=True)
MNIST(self.data_dir, train=False, download=True)

def setup(self, stage=None):
# Assign train/val datasets for use in dataloaders
if stage == "fit" or stage is None:
mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

# Assign test dataset for use in dataloader(s)
if stage == "test" or stage is None:
self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

def train_dataloader(self):
return DataLoader(self.mnist_train, batch_size=BATCH_SIZE)

def val_dataloader(self):
return DataLoader(self.mnist_val, batch_size=BATCH_SIZE)

def test_dataloader(self):
return DataLoader(self.mnist_test, batch_size=BATCH_SIZE)

if __name__ == '__main__':
model = LitMNIST()
checkpoint_callback = ModelCheckpoint(
dirpath='/Users/adam.amster/Downloads/pl_test',
monitor='val_loss',
save_top_k=2
)
trainer = L.Trainer(
accelerator="auto",
devices=1,
max_epochs=3,
default_root_dir='/Users/adam.amster/Downloads/pl_test',
callbacks=[checkpoint_callback]
)
# first fit, then comment, then run again
trainer.fit(model)
trainer.test(model=model, ckpt_path='best')
```

### Error messages and logs

```
`ValueError: `.test(ckpt_path="best")` is set but `ModelCheckpoint` is not configured to save the best model.`
```

### Environment

```

Current environment

* CUDA:
- GPU: None
- available: False
- version: None
* Lightning:
- lightning: 2.0.1
- lightning-cloud: 0.5.32
- lightning-utilities: 0.8.0
- pytorch-lightning: 2.0.1
- torch: 2.0.0
- torchmetrics: 0.11.4
- torchvision: 0.15.1
* Packages:
- aiohttp: 3.8.4
- aiosignal: 1.3.1
- anyio: 3.6.2
- arrow: 1.2.3
- async-timeout: 4.0.2
- attrs: 22.2.0
- beautifulsoup4: 4.12.2
- blessed: 1.20.0
- certifi: 2022.12.7
- charset-normalizer: 3.1.0
- click: 8.1.3
- contourpy: 1.0.7
- croniter: 1.3.10
- cycler: 0.11.0
- dateutils: 0.6.12
- deepdiff: 6.3.0
- dnspython: 2.3.0
- email-validator: 1.3.1
- fastapi: 0.88.0
- filelock: 3.11.0
- fonttools: 4.39.3
- frozenlist: 1.3.3
- fsspec: 2023.4.0
- h11: 0.14.0
- httpcore: 0.16.3
- httptools: 0.5.0
- httpx: 0.23.3
- idna: 3.4
- imageio: 2.27.0
- importlib-resources: 5.12.0
- inquirer: 3.1.3
- itsdangerous: 2.1.2
- jinja2: 3.1.2
- joblib: 1.2.0
- kiwisolver: 1.4.4
- lazy-loader: 0.2
- lightning: 2.0.1
- lightning-cloud: 0.5.32
- lightning-utilities: 0.8.0
- markdown-it-py: 2.2.0
- markupsafe: 2.1.2
- matplotlib: 3.7.1
- mdurl: 0.1.2
- mpmath: 1.3.0
- multidict: 6.0.4
- networkx: 3.0
- numpy: 1.24.2
- opencv-python: 4.7.0.72
- ordered-set: 4.1.0
- orjson: 3.8.10
- packaging: 23.0
- pandas: 1.5.3
- pillow: 9.5.0
- pip: 23.0.1
- psutil: 5.9.4
- pydantic: 1.10.7
- pygments: 2.14.0
- pyjwt: 2.6.0
- pyparsing: 3.0.9
- pytesseract: 0.3.10
- python-dateutil: 2.8.2
- python-dotenv: 1.0.0
- python-editor: 1.0.4
- python-multipart: 0.0.6
- pytorch-lightning: 2.0.1
- pytz: 2023.3
- pywavelets: 1.4.1
- pyyaml: 6.0
- readchar: 4.0.5
- requests: 2.28.2
- rfc3986: 1.5.0
- rich: 13.3.3
- scikit-image: 0.20.0
- scikit-learn: 1.2.2
- scipy: 1.9.1
- seaborn: 0.12.2
- setuptools: 67.4.0
- six: 1.16.0
- sniffio: 1.3.0
- soupsieve: 2.4
- starlette: 0.22.0
- starsessions: 1.3.0
- sympy: 1.11.1
- threadpoolctl: 3.1.0
- tifffile: 2023.3.21
- torch: 2.0.0
- torchmetrics: 0.11.4
- torchvision: 0.15.1
- tqdm: 4.65.0
- traitlets: 5.9.0
- typing-extensions: 4.5.0
- ujson: 5.7.0
- urllib3: 1.26.15
- uvicorn: 0.21.1
- uvloop: 0.17.0
- watchfiles: 0.19.0
- wcwidth: 0.2.6
- websocket-client: 1.5.1
- websockets: 11.0.1
- wheel: 0.38.4
- yarl: 1.8.2
- zipp: 3.15.0
* System:
- OS: Darwin
- architecture:
- 64bit
-
- processor: i386
- python: 3.8.2
- version: Darwin Kernel Version 21.6.0: Mon Aug 22 20:17:10 PDT 2022; root:xnu-8020.140.49~2/RELEASE_X86_64

```

### 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.test entry point and the ModelCheckpoint handling used by ckpt_path='best', using the reproduction in the issue to compare calls made after and without fit. Trace why the saved best checkpoint is not recognized, then verify that trainer.test(model=model, ckpt_path='best') works with the shown callback configuration.

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.