from collections.abc import MutableMapping
import pytorch_lightning as pl
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
[docs]
def flatten_dict(d, parent_key="", sep="_"):
"""Flattens a dictionary into a single-level dictionary while preserving
parent keys. Taken from
`SO <https://stackoverflow.com/questions/6027558/flatten-nested-dictionaries-compressing-keys>`_
Args:
d (MutableMapping): Dictionary to be flattened.
parent_key (str): String to use as a prefix to all subsequent keys.
sep (str): String to use as a separator between two key levels.
Returns:
dict: Single-level dictionary, flattened.
"""
items = []
for k, v in d.items():
new_key = parent_key + sep + k if parent_key else k
if isinstance(v, MutableMapping):
items.extend(flatten_dict(v, new_key, sep=sep).items())
else:
items.append((new_key, v))
return dict(items)
[docs]
class System(pl.LightningModule):
"""Base class for deep learning systems.
Contains a model, an optimizer, a loss function, training and validation
dataloaders and learning rate scheduler.
Note that by default, any PyTorch-Lightning hooks are *not* passed to the model.
If you want to use Lightning hooks, add the hooks to a subclass::
class MySystem(System):
def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
return self.model.on_train_batch_start(batch, batch_idx, dataloader_idx)
Args:
model (torch.nn.Module): Instance of model.
optimizer (torch.optim.Optimizer): Instance or list of optimizers.
loss_func (callable): Loss function with signature
(est_targets, targets).
train_loader (torch.utils.data.DataLoader): Training dataloader.
val_loader (torch.utils.data.DataLoader): Validation dataloader.
scheduler (torch.optim.lr_scheduler._LRScheduler): Instance, or list
of learning rate schedulers. Also supports dict or list of dict as
``{"interval": "step", "scheduler": sched}`` where ``interval=="step"``
for step-wise schedulers and ``interval=="epoch"`` for classical ones.
config: Anything to be saved with the checkpoints during training.
The config dictionary to re-instantiate the run for example.
.. note:: By default, ``training_step`` (used by ``pytorch-lightning`` in the
training loop) and ``validation_step`` (used for the validation loop)
share ``common_step``. If you want different behavior for the training
loop and the validation loop, overwrite both ``training_step`` and
``validation_step`` instead.
For more info on its methods, properties and hooks, have a look at lightning's docs:
https://pytorch-lightning.readthedocs.io/en/stable/lightning_module.html#lightningmodule-api
"""
default_monitor: str = "val_loss"
def __init__(
self,
model,
optimizer,
loss_func,
train_loader,
val_loader=None,
scheduler=None,
config=None,
):
super().__init__()
self.model = model
self.optimizer = optimizer
self.loss_func = loss_func
self.train_loader = train_loader
self.val_loader = val_loader
self.scheduler = scheduler
self.config = {} if config is None else config
# Save lightning's AttributeDict under self.hparams
self.save_hyperparameters(self.config_to_hparams(self.config))
[docs]
def forward(self, *args, **kwargs):
"""Applies forward pass of the model.
Returns:
:class:`torch.Tensor`
"""
return self.model(*args, **kwargs)
[docs]
def common_step(self, batch, batch_nb, train=True):
"""Common forward step between training and validation.
The function of this method is to unpack the data given by the loader,
forward the batch through the model and compute the loss.
Pytorch-lightning handles all the rest.
Args:
batch: the object returned by the loader (a list of torch.Tensor
in most cases) but can be something else.
batch_nb (int): The number of the batch in the epoch.
train (bool): Whether in training mode. Needed only if the training
and validation steps are fundamentally different, otherwise,
pytorch-lightning handles the usual differences.
Returns:
:class:`torch.Tensor` : The loss value on this batch.
.. note::
This is typically the method to overwrite when subclassing
``System``. If the training and validation steps are somehow
different (except for ``loss.backward()`` and ``optimzer.step()``),
the argument ``train`` can be used to switch behavior.
Otherwise, ``training_step`` and ``validation_step`` can be overwriten.
"""
inputs, targets = batch
est_targets = self(inputs)
loss = self.loss_func(est_targets, targets)
return loss
[docs]
def training_step(self, batch, batch_nb):
"""Pass data through the model and compute the loss.
Backprop is **not** performed (meaning PL will do it for you).
Args:
batch: the object returned by the loader (a list of torch.Tensor
in most cases) but can be something else.
batch_nb (int): The number of the batch in the epoch.
Returns:
torch.Tensor, the value of the loss.
"""
loss = self.common_step(batch, batch_nb, train=True)
self.log("loss", loss, logger=True)
return loss
[docs]
def validation_step(self, batch, batch_nb):
"""Need to overwrite PL validation_step to do validation.
Args:
batch: the object returned by the loader (a list of torch.Tensor
in most cases) but can be something else.
batch_nb (int): The number of the batch in the epoch.
"""
loss = self.common_step(batch, batch_nb, train=False)
self.log("val_loss", loss, on_epoch=True, prog_bar=True)
[docs]
def on_validation_epoch_end(self):
"""Log hp_metric to tensorboard for hparams selection."""
hp_metric = self.trainer.callback_metrics.get("val_loss", None)
if hp_metric is not None:
self.trainer.logger.log_metrics(
{"hp_metric": hp_metric}, step=self.trainer.global_step
)
[docs]
def lr_scheduler_step(self, scheduler, metric):
if metric is None:
scheduler.step()
else:
scheduler.step(metric)
[docs]
def train_dataloader(self):
"""Training dataloader"""
return self.train_loader
[docs]
def val_dataloader(self):
"""Validation dataloader"""
return self.val_loader
[docs]
def on_save_checkpoint(self, checkpoint):
"""Overwrite if you want to save more things in the checkpoint."""
checkpoint["training_config"] = self.config
return checkpoint
[docs]
@staticmethod
def config_to_hparams(dic):
"""Sanitizes the config dict to be handled correctly by torch
SummaryWriter. It flatten the config dict, converts ``None`` to
``"None"`` and any list and tuple into torch.Tensors.
Args:
dic (dict): Dictionary to be transformed.
Returns:
dict: Transformed dictionary.
"""
dic = flatten_dict(dic)
for k, v in dic.items():
if v is None:
dic[k] = str(v)
elif isinstance(v, (list, tuple)):
dic[k] = torch.tensor(v)
return dic