import uuid import torch from pytorch_lightning import Trainer, LightningModule from pytorch_lightning.callbacks import ModelCheckpoint from torch.utils.data import DataLoader, Dataset class RandomDataset(Dataset): def __init__(self, size, length): self.len = length self.data = torch.randn(length, size) def __getitem__(self, index): return self.data[index] def __len__(self): return self.len class BoringModel(LightningModule): def __init__(self): super().__init__() self.layer = torch.nn.Linear(32, 2) def forward(self, x): return self.layer(x) def training_step(self, batch, batch_idx): loss = self(batch).sum() self.log("train_loss", loss) return {"loss": loss} def validation_step(self, batch, batch_idx): loss = self(batch).sum() self.log("valid_loss", loss) def test_step(self, batch, batch_idx): loss = self(batch).sum() self.log("test_loss", loss) def configure_optimizers(self): return torch.optim.SGD(self.layer.parameters(), lr=0.1) if __name__ == "__main__": tmpdir = f"/tmp/{uuid.uuid4()}" print(tmpdir) trainer = Trainer( default_root_dir=tmpdir, max_epochs=2, callbacks=[ModelCheckpoint(dirpath=tmpdir, every_n_epochs=10, save_last='link', save_top_k=-1)], enable_checkpointing=True, ) model = BoringModel() trainer.fit(model, train_dataloaders=DataLoader(RandomDataset(32, 64), batch_size=2))