File size: 1,443 Bytes
9f31860 2bb6467 9f31860 2bb6467 9f31860 2bb6467 9f31860 2bb6467 9f31860 2bb6467 9f31860 30df46a 2bb6467 9f31860 2bb6467 9f31860 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import lightning as L
import torch
from lightning.pytorch.callbacks import (
ModelCheckpoint,
LearningRateMonitor,
EarlyStopping,
)
from lightning.pytorch.loggers import TensorBoardLogger
from src.dataset import DRDataModule
from src.model import DRModel
# seed everything for reproducibility
SEED = 42
L.seed_everything(SEED, workers=True)
torch.set_float32_matmul_precision("high")
# Init DataModule
dm = DRDataModule(batch_size=128, num_workers=24)
dm.setup()
# Init model from datamodule's attributes
model = DRModel(
num_classes=dm.num_classes, learning_rate=3e-4, class_weights=dm.class_weights
)
# Init logger
logger = TensorBoardLogger(save_dir="artifacts")
# Init callbacks
checkpoint_callback = ModelCheckpoint(
monitor="val_loss",
mode="min",
save_top_k=2,
dirpath="artifacts/checkpoints",
filename="{epoch}-{step}-{val_loss:.2f}-{val_acc:.2f}-{val_kappa:.2f}",
)
# Init LearningRateMonitor
lr_monitor = LearningRateMonitor(logging_interval="step")
# early stopping
early_stopping = EarlyStopping(
monitor="val_loss",
patience=5,
verbose=True,
mode="min",
)
# Init trainer
trainer = L.Trainer(
max_epochs=20,
accelerator="auto",
devices="auto",
logger=logger,
callbacks=[checkpoint_callback, lr_monitor, early_stopping],
# check_val_every_n_epoch=4,
)
# Pass the datamodule as arg to trainer.fit to override model hooks :)
trainer.fit(model, dm)
|