File size: 3,648 Bytes
cbe6208 |
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 |
"""Custom callbacks
"""
from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import BaseFinetuning, EarlyStopping, LearningRateFinder
from lightning.pytorch.trainer.states import TrainerFn
class PhaseEarlyStopping(EarlyStopping):
"""Monitor a validation metric and stop training when it stops improving.
Only functions in a specific phase of training.
"""
training_phase = None
def switch_phase(self, phase: str):
"""Switch phase of callback"""
if phase == self.training_phase:
self.activate()
else:
self.deactivate()
def deactivate(self):
"""Deactivate callback"""
self.active = False
def activate(self):
"""Activate callback"""
self.active = True
def _should_skip_check(self, trainer: Trainer) -> bool:
return (
(trainer.state.fn != TrainerFn.FITTING) or (trainer.sanity_checking) or not self.active
)
class PretrainEarlyStopping(EarlyStopping):
"""Monitor a validation metric and stop training when it stops improving.
Only functions in the 'pretrain' phase of training.
"""
training_phase = "pretrain"
class MainEarlyStopping(EarlyStopping):
"""Monitor a validation metric and stop training when it stops improving.
Only functions in the 'main' phase of training.
"""
training_phase = "main"
class PretrainFreeze(BaseFinetuning):
"""Freeze the satellite and NWP encoders during pretraining"""
training_phase = "pretrain"
def __init__(self):
"""Freeze the satellite and NWP encoders during pretraining"""
super().__init__()
def freeze_before_training(self, pl_module):
"""Freeze satellite and NWP encoders before training start"""
# freeze any module you want
modules = []
if pl_module.include_sat:
modules += [pl_module.sat_encoder]
if pl_module.include_nwp:
modules += [pl_module.nwp_encoder]
self.freeze(modules)
def finetune_function(self, pl_module, current_epoch, optimizer):
"""Unfreeze satellite and NWP encoders"""
if not self.active:
modules = []
if pl_module.include_sat:
modules += [pl_module.sat_encoder]
if pl_module.include_nwp:
modules += [pl_module.nwp_encoder]
self.unfreeze_and_add_param_group(
modules=modules,
optimizer=optimizer,
train_bn=True,
)
def switch_phase(self, phase: str):
"""Switch phase of callback"""
if phase == self.training_phase:
self.activate()
else:
self.deactivate()
def deactivate(self):
"""Deactivate callback"""
self.active = False
def activate(self):
"""Activate callback"""
self.active = True
class PhasedLearningRateFinder(LearningRateFinder):
"""Finds a learning rate at the start of each phase of learning"""
active = True
def on_fit_start(self, *args, **kwargs):
"""Do nothing"""
return
def on_train_epoch_start(self, trainer, pl_module):
"""Run learning rate finder on epoch start and then deactivate"""
if self.active:
self.lr_find(trainer, pl_module)
self.deactivate()
def switch_phase(self, phase: str):
"""Switch training phase"""
self.activate()
def deactivate(self):
"""Deactivate callback"""
self.active = False
def activate(self):
"""Activate callback"""
self.active = True
|