pytorch-lightning
pytorch-lightning copied to clipboard
Warning is raised in PyTorch 2.2 and Lightning 2.2 when using the Tuner
Bug description
A user reported that there is now a user warning when using the Tuner
UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`
What version are you seeing the problem on?
v2.2
How to reproduce the bug
`
import torch
class PyTorchMLP(torch.nn.Module):
def __init__(self, num_features, num_classes):
super().__init__()
self.all_layers = torch.nn.Sequential(
# 1st hidden layer
torch.nn.Linear(num_features, 100),
torch.nn.ReLU(),
# 2nd hidden layer
torch.nn.Linear(100, 50),
torch.nn.ReLU(),
# output layer
torch.nn.Linear(50, num_classes),
)
def forward(self, x):
x = torch.flatten(x, start_dim=1)
logits = self.all_layers(x)
return logits
class LightningModel(L.LightningModule):
def __init__(self, model, learning_rate):
super().__init__()
self.learning_rate = learning_rate
self.model = model
self.save_hyperparameters(ignore=["model"])
self.train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=2)
self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=2)
self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=2)
def forward(self, x):
return self.model(x)
def _shared_step(self, batch):
features, true_labels = batch
logits = self(features)
loss = F.cross_entropy(logits, true_labels)
predicted_labels = torch.argmax(logits, dim=1)
return loss, true_labels, predicted_labels
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("train_loss", loss)
self.train_acc(predicted_labels, true_labels)
self.log(
"train_acc", self.train_acc, prog_bar=True, on_epoch=True, on_step=False
)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("val_loss", loss, prog_bar=True)
self.val_acc(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, prog_bar=True)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels = self._shared_step(batch)
self.test_acc(predicted_labels, true_labels)
self.log("test_acc", self.test_acc)
def configure_optimizers(self):
optimizer = torch.optim.SGD(self.parameters(), lr=self.learning_rate)
return optimizer
class CustomDataset(Dataset):
def __init__(self, feature_array, label_array, transform=None):
self.x = feature_array
self.y = label_array
self.transform = transform
def __getitem__(self, index):
x = self.x[index]
y = self.y[index]
if self.transform is not None:
x = self.transform(x)
return x, y
def __len__(self):
return self.y.shape[0]
class CustomDataModule(L.LightningDataModule):
def __init__(self, data_dir="./mnist", batch_size=64):
super().__init__()
self.data_dir = data_dir
self.batch_size = batch_size
def prepare_data(self):
# download
pass
def setup(self, stage: str):
X, y = make_classification(
n_samples=20000,
n_features=100,
n_informative=10,
n_redundant=40,
n_repeated=25,
n_clusters_per_class=5,
flip_y=0.05,
class_sep=0.5,
random_state=123,
)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=1, shuffle=True
)
X_train, X_val, y_train, y_val = train_test_split(
X_train, y_train, test_size=0.1, random_state=1, shuffle=True
)
self.train_dataset = CustomDataset(
feature_array=X_train.astype(np.float32),
label_array=y_train.astype(np.int64),
)
self.val_dataset = CustomDataset(
feature_array=X_val.astype(np.float32), label_array=y_val.astype(np.int64)
)
self.test_dataset = CustomDataset(
feature_array=X_test.astype(np.float32), label_array=y_test.astype(np.int64)
)
def train_dataloader(self):
train_loader = DataLoader(
dataset=self.train_dataset,
batch_size=32,
shuffle=True,
drop_last=True,
num_workers=0,
)
return train_loader
def val_dataloader(self):
val_loader = DataLoader(
dataset=self.val_dataset,
batch_size=32,
shuffle=False,
num_workers=0,
)
return val_loader
def test_dataloader(self):
test_loader = DataLoader(
dataset=self.test_dataset, batch_size=32, shuffle=False, num_workers=0
)
return test_loader
from pytorch_lightning.tuner import Tuner
trainer = L.Trainer(
max_epochs=100,
)
torch.manual_seed(123)
dm = CustomDataModule()
pytorch_model = PyTorchMLP(num_features=100, num_classes=2)
lightning_model = LightningModel(model=pytorch_model, learning_rate=0.1)
# Create a Tuner
tuner = Tuner(trainer)
lr_finder = tuner.lr_find(lightning_model, datamodule=dm)
### Error messages and logs
Error messages and logs here please
### Environment
<details>
<summary>Current environment</summary>
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(conda, pip, source):
#- Running environment of LightningApp (e.g. local, cloud):
</details>
### More info
_No response_
Hi @rasbt
The code above is missing these imports:
from sklearn.datasets import make_classification
from torch.utils.data import DataLoader, Dataset
import numpy as np
from shared_utilities import train_test_split
It doesn't generate any scheduler warnings for me.
Then I pulled all notebooks from https://github.com/Lightning-AI/dl-fundamentals/tree/main/unit06-dl-tips/6.2-learning-rates, installed the dependencies (there is no requirements.txt file), and ran the notebooks. There are no warnings regarding the scheduler in any of the cell outputs.
Ah sorry for the incomplete code. Yeah, it seems weird, I am not getting any warnings, either. I think we can close this and sorry for bothering with that!