Spaces:
Sleeping
Sleeping
File size: 2,552 Bytes
7435e6e | 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 | import pytorch_lightning as pl
import torch
from hydra.utils import instantiate
class AbstractModel(pl.LightningModule):
def __init__(self,
lr=0.001,
optimizer_hparams=dict(),
scheduler=dict(classname='MultiStepLR', kwargs=dict(milestones=[100, 150], gamma=0.1))
):
super().__init__()
# Exports the hyperparameters to a YAML file, and create "self.hparams" namespace
self.save_hyperparameters()
def forward(self, x):
raise NotImplementedError("Subclass needs to implement this method")
def configure_optimizers(self):
# AdamW is Adam with a correct implementation of weight decay (see here
# for details: https://arxiv.org/pdf/1711.05101.pdf)
print("configuring the optimizer and lr scheduler with learning rate=%.5f"%self.hparams.lr)
optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr, **self.hparams.optimizer_hparams)
# scheduler = getattr(torch.optim.lr_scheduler, self.hparams.lr_hparams['classname'])(optimizer, **self.hparams.lr_hparams['kwargs'])
if self.hparams.scheduler is not None:
scheduler = instantiate({**self.hparams.scheduler, '_partial_': True})(optimizer)
return [optimizer], [scheduler]
else:
return optimizer
def additional_losses(self):
"""get additional_losses"""
return torch.zeros((1))
def process_batch_supervised(self, batch):
"""get predictions, losses and mean errors (MAE)"""
raise NotImplementedError("Subclass needs to implement this method")
def log_all(self, losses, metrics, prefix=''):
for k, v in losses.items():
self.log(f'{prefix}{k}_loss', v.item() if isinstance(v, torch.Tensor) else v)
for k, v in metrics.items():
self.log(f'{prefix}{k}', v.item() if isinstance(v, torch.Tensor) else v)
def training_step(self, batch, batch_idx):
# "batch" is the output of the training data loader.
preds, losses, metrics = self.process_batch_supervised(batch)
self.log_all(losses, metrics, prefix='train_')
return losses['final']
def validation_step(self, batch, batch_idx):
preds, losses, metrics = self.process_batch_supervised(batch)
self.log_all(losses, metrics, prefix='val_')
def test_step(self, batch, batch_idx):
preds, losses, metrics = self.process_batch_supervised(batch)
self.log_all(losses, metrics, prefix='test_') |