| | import torch |
| | import pytorch_lightning as pl |
| | import torch.nn as nn |
| | import os |
| | import torch.optim.lr_scheduler as lrs |
| | import inspect |
| |
|
| |
|
| | class MInterface_base(pl.LightningModule): |
| | def __init__(self, model_name=None, loss=None, lr=None, **kargs): |
| | super().__init__() |
| | self.save_hyperparameters() |
| | self.load_model() |
| | self.configure_loss() |
| | os.makedirs(os.path.join(self.hparams.res_dir, self.hparams.ex_name), exist_ok=True) |
| | |
| | def forward(self, input): |
| | pass |
| | |
| | |
| | def training_step(self, batch, batch_idx, **kwargs): |
| | pass |
| |
|
| |
|
| | def validation_step(self, batch, batch_idx): |
| | pass |
| |
|
| | def test_step(self, batch, batch_idx): |
| | |
| | return self.validation_step(batch, batch_idx) |
| |
|
| | def on_validation_epoch_end(self): |
| | |
| | self.print('') |
| | |
| | def get_schedular(self, optimizer, lr_scheduler='onecycle'): |
| | if lr_scheduler == 'step': |
| | scheduler = lrs.StepLR(optimizer, |
| | step_size=self.hparams.lr_decay_steps, |
| | gamma=self.hparams.lr_decay_rate) |
| | elif lr_scheduler == 'cosine': |
| | scheduler = lrs.CosineAnnealingLR(optimizer, |
| | T_max=self.hparams.lr_decay_steps) |
| | elif lr_scheduler == 'onecycle': |
| | scheduler = lrs.OneCycleLR(optimizer, max_lr=self.hparams.lr, T_max=self.hparams.lr_decay_steps, three_phase=False) |
| | elif lr_scheduler == 'plateau': |
| | scheduler = lrs.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3, verbose=True) |
| | else: |
| | raise ValueError('Invalid lr_scheduler type!') |
| |
|
| | return scheduler |
| |
|
| | def configure_optimizers(self): |
| | if hasattr(self.hparams, 'weight_decay'): |
| | weight_decay = self.hparams.weight_decay |
| | else: |
| | weight_decay = 0 |
| | |
| | optimizer_g = torch.optim.AdamW(filter(lambda p: p.requires_grad, self.model.parameters()), lr=self.hparams.lr, weight_decay=weight_decay, betas=(0.9, 0.98), eps=1e-8) |
| |
|
| |
|
| | schecular_g = self.get_schedular(optimizer_g, self.hparams.lr_scheduler) |
| | if self.hparams.lr_scheduler == 'plateau': |
| | return [optimizer_g], [{"scheduler": schecular_g, "interval": "epoch", "monitor": "val_loss"}] |
| | else: |
| | return [optimizer_g], [{"scheduler": schecular_g, "interval": "step"}] |
| | |
| | def lr_scheduler_step(self, *args, **kwargs): |
| | scheduler = self.lr_schedulers() |
| | if self.hparams.lr_scheduler != 'plateau': |
| | scheduler.step() |
| | |
| | |
| | def configure_devices(self): |
| | self.device = torch.device(self.hparams.device) |
| |
|
| | def configure_loss(self): |
| | self.loss_function = nn.CrossEntropyLoss(reduction='none') |
| | |
| | def load_model(self): |
| | self.model = None |
| |
|
| | def instancialize(self, Model, **other_args): |
| | """ Instancialize a model using the corresponding parameters |
| | from self.hparams dictionary. You can also input any args |
| | to overwrite the corresponding value in self.hparams. |
| | """ |
| | class_args = inspect.getargspec(Model.__init__).args[1:] |
| | inkeys = self.hparams.keys() |
| | args1 = {} |
| | for arg in class_args: |
| | if arg in inkeys: |
| | args1[arg] = getattr(self.hparams, arg) |
| | args1.update(other_args) |
| | return Model(**args1) |
| |
|
