Lightning-AI / Lightning-AI/pytorch-lightning
trainer.test(ckpt_path='best') does not work as expected
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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