From 6c8cffe28baeafac77d047ff3f8ded9418033e2f Mon Sep 17 00:00:00 2001 From: Volpeon Date: Mon, 16 Jan 2023 15:52:43 +0100 Subject: More training adjustments --- training/optimization.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) (limited to 'training/optimization.py') diff --git a/training/optimization.py b/training/optimization.py index 5db7794..6dee4bc 100644 --- a/training/optimization.py +++ b/training/optimization.py @@ -49,8 +49,8 @@ def get_one_cycle_schedule( annealing: Literal["cos", "half_cos", "linear"] = "cos", warmup_exp: int = 1, annealing_exp: int = 1, - min_lr: int = 0.04, - mid_point: int = 0.3, + min_lr: float = 0.04, + mid_point: float = 0.3, last_epoch: int = -1 ): if warmup == "linear": @@ -91,10 +91,10 @@ def get_scheduler( id: str, optimizer: torch.optim.Optimizer, num_training_steps_per_epoch: int, - gradient_accumulation_steps: int, + gradient_accumulation_steps: int = 1, min_lr: float = 0.04, - warmup_func: str = "cos", - annealing_func: str = "cos", + warmup_func: Literal["cos", "linear"] = "cos", + annealing_func: Literal["cos", "half_cos", "linear"] = "cos", warmup_exp: int = 1, annealing_exp: int = 1, cycles: int = 1, -- cgit v1.2.3-54-g00ecf