File size: 7,974 Bytes
d9115cb | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | from typing import Any, Dict, Optional, Union
import torch
import transformers
from lightning import LightningModule
from torchmetrics import MeanMetric
from torchmetrics.aggregation import BaseAggregator
from src.components.eval_metrics import Evaluator
from src.utils.pylogger import RankedLogger
command_line_logger = RankedLogger(__name__, rank_zero_only=True)
class BaseModule(LightningModule):
def __init__(
self,
model: Union[torch.nn.Module, transformers.PreTrainedModel],
optimizer: torch.optim.Optimizer,
scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
loss_function: torch.nn.Module,
evaluator: Evaluator,
training_loop_function: callable = None,
) -> None:
"""
Args:
model: The model to train.
optimizer: The optimizer to use for the model.
scheduler: The scheduler to use for the model.
loss_function: The loss function to use for the model.
evaluator: The evaluator to use for the model.
training_loop_function: The training loop function to use for the model, in case it is different than the default one.
"""
super().__init__()
self.model = model
self.optimizer = optimizer
self.scheduler = scheduler
self.loss_function = loss_function
self.evaluator = evaluator
self.training_loop_function = training_loop_function
# We use setters to set the prediction key and name.
self._prediction_key_name = None
self._prediction_name = None
if self.training_loop_function is not None:
self.automatic_optimization = False
if self.evaluator: # For inference, evaluator is not set.
for metric_name, metric_object in self.evaluator.metrics.items():
setattr(self, metric_name, metric_object)
# for averaging loss across batches
self.train_loss = MeanMetric()
self.val_loss = MeanMetric()
self.test_loss = MeanMetric()
@property
def prediction_key_name(self) -> Optional[str]:
return self._prediction_key_name
@prediction_key_name.setter
def prediction_key_name(self, value: str) -> None:
command_line_logger.debug(f"Setting prediction_key_name to {value}")
self._prediction_key_name = value
@property
def prediction_name(self) -> Optional[str]:
return self._prediction_name
@prediction_name.setter
def prediction_name(self, value: str) -> None:
command_line_logger.debug(f"Setting prediction_name to {value}")
self._prediction_name = value
def forward(
self,
**kwargs: Dict[str, torch.Tensor],
) -> torch.Tensor:
raise NotImplementedError(
"Inherit from this class and implement the forward method."
)
def model_step(
self,
model_input: Any,
label_data: Optional[Any] = None,
):
raise NotImplementedError(
"Inherit from this class and implement the model_step method."
)
def on_train_start(self) -> None:
"""Lightning hook that is called when training begins."""
# by default lightning executes validation step sanity checks before training starts,
# so it's worth to make sure validation metrics don't store results from these checks
self.val_loss.reset()
self.evaluator.reset()
self.train_loss.reset()
self.test_loss.reset()
def on_validation_epoch_start(self) -> None:
"""Lightning hook that is called when a validation epoch starts."""
self.val_loss.reset()
self.evaluator.reset()
def on_test_epoch_start(self):
self.test_loss.reset()
self.evaluator.reset()
def on_validation_epoch_end(self) -> None:
"Lightning hook that is called when a validation epoch ends."
self.log("val/loss", self.val_loss, sync_dist=False, prog_bar=True, logger=True)
self.log_metrics("val")
def on_test_epoch_end(self) -> None:
self.log(
"test/loss", self.test_loss, sync_dist=False, prog_bar=True, logger=True
)
self.log_metrics("test")
def on_exception(self, exception):
self.trainer.should_stop = True # stop all workers
self.trainer.logger.finalize(status="failure")
def log_metrics(
self,
prefix: str,
on_step=False,
on_epoch=True,
# We use sync_dist=False by default because, if using retrieval metrics, those are already synchronized. Change if using
# different metrics than the default ones.
sync_dist=False,
logger=True,
prog_bar=False,
call_compute=False,
) -> Dict[str, Any]:
metrics_dict = {
f"{prefix}/{metric_name}": metric_object.compute()
if call_compute
else metric_object
for metric_name, metric_object in self.evaluator.metrics.items()
}
self.log_dict(
metrics_dict,
on_step=on_step,
on_epoch=on_epoch,
sync_dist=sync_dist,
logger=logger,
prog_bar=prog_bar,
)
def setup(self, stage: str) -> None:
"""Lightning hook that is called at the beginning of fit (train + validate), validate,
test, or predict.
This is a good hook when you need to build models dynamically or adjust something about
them. This hook is called on every process when using DDP.
:param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
"""
# if self.hparams.compile and stage == "fit":
# self.net = torch.compile(self.net)
pass
def configure_optimizers(self) -> Dict[str, Any]:
"""Choose what optimizers and learning-rate schedulers to use in your optimization.
Normally you'd need one. But in the case of GANs or similar you might have multiple.
Examples:
https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
:return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
"""
optimizer = self.optimizer(params=self.trainer.model.parameters())
if self.scheduler is not None:
scheduler = self.scheduler(optimizer=optimizer)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val/loss",
"interval": "step",
"frequency": 1,
},
}
return {"optimizer": optimizer}
def eval_step(self, batch: Any, loss_to_aggregate: BaseAggregator):
raise NotImplementedError("eval_step method must be implemented.")
def validation_step(
self,
batch: Any,
batch_idx: int,
) -> None:
"""Perform a single validation step on a batch of data from the validation set.
:param batch: A batch of data of data (tuple) where first object is a SequentialModelInputData object
and second is a SequentialModuleLabelData object.
"""
self.eval_step(batch, self.val_loss)
def test_step(
self,
batch: Any,
batch_idx: int,
) -> None:
"""Perform a single test step on a batch of data from the test set.
:param batch: A batch of data of data (tuple) where first object is a SequentialModelInputData object
and second is a SequentialModuleLabelData object.
"""
self.eval_step(batch, self.test_loss)
|