awaelchli
commented
11 months ago @jscottcronin All your methods are indented under the __init__ method. You should unindent them one tab. Here's the fixed code:
from typing import Any
import lightning.pytorch as pl
from lightning.pytorch.utilities.types import STEP_OUTPUT, OptimizerLRScheduler
from torch import optim, nn
import torch.nn.functional as F
import torch
from torch.utils.data import random_split, DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from typing import Optional
class MNISTModel(pl.LightningModule):
def __init__(self, learning_rate=1e-3):
super().__init__()
self.learning_rate = learning_rate
self.model = nn.Sequential(
nn.Flatten(),
nn.Linear(28 * 28, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(128, 10),
nn.ReLU(),
)
def cross_entropy_loss(self, logits, labels):
return F.nll_loss(logits, labels)
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
# ys = torch.zeros(x.size(0), 10)
# ys[y] = 1
logits = self.forward(x)
loss = self.cross_entropy_loss(logits, y)
self.log("train_loss", loss)
return loss
def validation_step(self, batch, batch_idx):
x, y = batch
# ys = torch.zeros(x.size(0), 10)
# ys[y] = 1
logits = self.forward(x)
loss = self.cross_entropy_loss(logits, y)
self.log("val_loss", loss)
return loss
def configure_optimizers(self) -> OptimizerLRScheduler:
optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
return optimizer
class MNISTDataModule(pl.LightningDataModule):
def __init__(self, batch_size=32) -> None:
super().__init__()
self.data_dir = "./data"
self.num_classes = 10
self.dims = (1, 28, 28)
self.batch_size = batch_size
self.transform = transforms.Compose(
[transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
)
def prepare_data(self) -> None:
MNIST(root="./data", train=True, transform=self.transform, download=True)
MNIST(root="./data", train=False, transform=self.transform, download=True)
def setup(self, stage: Optional[str] = None) -> None:
if stage == "fit" or stage is None:
full = MNIST(root=self.data_dir, train=True, transform=self.transform)
self.train, self.val = random_split(full, [55000, 5000])
if stage == "test" or stage is None:
self.test = MNIST(
root=self.data_dir, train=False, transform=self.transform
)
def train_dataloader(self) -> DataLoader:
return DataLoader(self.train, batch_size=self.batch_size)
def val_dataloader(self) -> DataLoader:
return DataLoader(self.val, batch_size=self.batch_size)
def test_dataloader(self) -> DataLoader:
return DataLoader(self.test, batch_size=self.batch_size)
if __name__ == "__main__":
model = MNISTModel()
dm = MNISTDataModule()
trainer = pl.Trainer(fast_dev_run=True, max_epochs=10)
trainer.fit(model, datamodule=dm)
# trainer.test(model, datamodule=dm)
I'm closing the issue because I don't see any action items. The error message was correct.
Bug description
Following the most basic example for MNIST with a multilayer model. Does not run properly on my Mac M1 laptop.
What version are you seeing the problem on?
v2.0
How to reproduce the bug
Error messages and logs
Environment
More info
No response