Final Train Runs

This commit is contained in:
Steffen Illium
2021-03-18 07:45:07 +01:00
parent ad254dae92
commit fecf4923c2
14 changed files with 672 additions and 362 deletions

View File

@@ -24,11 +24,16 @@ class OptimizerMixin:
if self.params.sto_weight_avg:
optimizer = SWA(optimizer, swa_start=10, swa_freq=5, swa_lr=0.05)
optimizer_dict.update(optimizer=optimizer)
if self.params.lr_warm_restart_epochs:
scheduler = CosineAnnealingWarmRestarts(optimizer, self.params.lr_warm_restart_epochs)
if self.params.scheduler == CosineAnnealingWarmRestarts.__name__:
scheduler = CosineAnnealingWarmRestarts(optimizer, self.params.lr_scheduler_parameter)
elif self.params.scheduler == LambdaLR.__name__:
lr_reduce_ratio = self.params.lr_scheduler_parameter
scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: lr_reduce_ratio ** epoch)
else:
scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: 0.95 ** epoch)
optimizer_dict.update(lr_scheduler=scheduler)
scheduler = None
if scheduler:
optimizer_dict.update(lr_scheduler=scheduler)
return optimizer_dict
def on_train_end(self):