diff options
Diffstat (limited to 'src/training/trainer/callbacks/__init__.py')
| -rw-r--r-- | src/training/trainer/callbacks/__init__.py | 15 |
1 files changed, 13 insertions, 2 deletions
diff --git a/src/training/trainer/callbacks/__init__.py b/src/training/trainer/callbacks/__init__.py index 5942276..c81e4bf 100644 --- a/src/training/trainer/callbacks/__init__.py +++ b/src/training/trainer/callbacks/__init__.py @@ -1,7 +1,16 @@ """The callback modules used in the training script.""" -from .base import Callback, CallbackList, Checkpoint +from .base import Callback, CallbackList +from .checkpoint import Checkpoint from .early_stopping import EarlyStopping -from .lr_schedulers import CyclicLR, MultiStepLR, OneCycleLR, ReduceLROnPlateau, StepLR +from .lr_schedulers import ( + CosineAnnealingLR, + CyclicLR, + MultiStepLR, + OneCycleLR, + ReduceLROnPlateau, + StepLR, + SWA, +) from .progress_bar import ProgressBar from .wandb_callbacks import WandbCallback, WandbImageLogger @@ -9,6 +18,7 @@ __all__ = [ "Callback", "CallbackList", "Checkpoint", + "CosineAnnealingLR", "EarlyStopping", "WandbCallback", "WandbImageLogger", @@ -18,4 +28,5 @@ __all__ = [ "ProgressBar", "ReduceLROnPlateau", "StepLR", + "SWA", ] |