166 lines
4.8 KiB
Python
166 lines
4.8 KiB
Python
import time
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import pytorch_lightning as pl
|
|
|
|
from torch import optim
|
|
from torchmetrics import MetricCollection
|
|
from torchmetrics.classification import Accuracy
|
|
|
|
|
|
class ClassificationLightning(pl.LightningModule):
|
|
def __init__(self,
|
|
model: nn.Module,
|
|
num_classes: int,
|
|
optimizer_class,
|
|
optimizer_kwargs: dict,
|
|
scheduler_class=None,
|
|
scheduler_kwargs: dict | None = None,
|
|
scheduler_monitor: str = "val/loss"):
|
|
super().__init__()
|
|
|
|
|
|
scheduler_kwargs = scheduler_kwargs or {}
|
|
|
|
self.model = model
|
|
self.optimizer_class = optimizer_class
|
|
self.scheduler_class = scheduler_class
|
|
|
|
num_parameters = sum(
|
|
p.numel()
|
|
for p in model.parameters()
|
|
)
|
|
|
|
num_trainable_parameters = sum(
|
|
p.numel()
|
|
for p in model.parameters()
|
|
if p.requires_grad
|
|
)
|
|
|
|
self.save_hyperparameters({
|
|
"model": model.__class__.__name__,
|
|
"num_classes": num_classes,
|
|
|
|
"num_parameters": num_parameters,
|
|
"num_trainable_parameters": num_trainable_parameters,
|
|
|
|
"optimizer": optimizer_class.__name__,
|
|
"optimizer_kwargs": optimizer_kwargs,
|
|
|
|
"scheduler": (
|
|
scheduler_class.__name__
|
|
if scheduler_class is not None
|
|
else None
|
|
),
|
|
"scheduler_kwargs": scheduler_kwargs,
|
|
"scheduler_monitor": scheduler_monitor
|
|
})
|
|
|
|
|
|
self.criterion = nn.CrossEntropyLoss()
|
|
self.train_metrics = MetricCollection({
|
|
"acc_top1": Accuracy(task="multiclass",
|
|
num_classes=self.hparams.num_classes,
|
|
top_k=1),
|
|
"acc_top5": Accuracy(task="multiclass",
|
|
num_classes=self.hparams.num_classes,
|
|
top_k=5),
|
|
})
|
|
|
|
self.val_metrics = self.train_metrics.clone()
|
|
self.test_metrics = self.train_metrics.clone()
|
|
|
|
def forward(self, image):
|
|
return self.model(image)
|
|
|
|
def _shared_step(self,
|
|
sample,
|
|
stage,
|
|
metrics):
|
|
x, y = sample
|
|
|
|
logits = self(x)
|
|
loss = self.criterion(logits, y)
|
|
|
|
metrics.update(logits, y)
|
|
|
|
self.log(f"{stage}/loss",
|
|
loss,
|
|
prog_bar=True,
|
|
on_epoch=True,
|
|
on_step=False
|
|
)
|
|
|
|
return loss
|
|
|
|
def _log_metrics(self, stage, metrics):
|
|
computed = metrics.compute()
|
|
|
|
for name, value in computed.items():
|
|
self.log(f"{stage}/{name}", value, prog_bar=True)
|
|
|
|
metrics.reset()
|
|
|
|
def on_train_epoch_start(self):
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
|
|
self._epoch_start_time = time.perf_counter()
|
|
|
|
def training_step(self, batch, batch_idx):
|
|
return self._shared_step(batch, "train", self.train_metrics)
|
|
|
|
def on_train_epoch_end(self):
|
|
self._log_metrics("train", self.train_metrics)
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
|
|
epoch_time = time.perf_counter() - self._epoch_start_time
|
|
lr = self.optimizers().param_groups[0]["lr"]
|
|
|
|
self.log("train/lr", lr, prog_bar=True)
|
|
self.log("train/epoch_time", epoch_time)
|
|
|
|
def validation_step(self, batch, batch_idx):
|
|
return self._shared_step(batch, "val", self.val_metrics)
|
|
|
|
def on_validation_epoch_end(self):
|
|
self._log_metrics("val", self.val_metrics)
|
|
|
|
def test_step(self, batch, batch_idx):
|
|
return self._shared_step(batch, "test", self.test_metrics)
|
|
|
|
def on_test_epoch_end(self):
|
|
self._log_metrics("test", self.test_metrics)
|
|
|
|
def configure_optimizers(self):
|
|
optimizer = self.optimizer_class(
|
|
self.parameters(),
|
|
**self.hparams.optimizer_kwargs
|
|
)
|
|
|
|
if self.scheduler_class is None:
|
|
return optimizer
|
|
|
|
scheduler = self.scheduler_class(
|
|
optimizer,
|
|
**self.hparams.scheduler_kwargs
|
|
)
|
|
|
|
scheduler_config = {
|
|
"scheduler": scheduler,
|
|
"interval": "epoch",
|
|
"frequency": 1
|
|
}
|
|
|
|
if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
|
|
scheduler_config["monitor"] = self.hparams.scheduler_monitor
|
|
|
|
return {
|
|
"optimizer": optimizer,
|
|
"lr_scheduler": scheduler_config
|
|
}
|
|
|