Final Train Runs
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user