10fromtypingimportTuple1112importtorch1314fromlabml.configsimportBaseConfigs,option,meta\_config15fromlabml\_nn.optimizersimportWeightDecay

18classOptimizerConfigs(BaseConfigs):

26optimizer:torch.optim.Adam

29weight\_decay\_obj:WeightDecay

32weight\_decouple:bool=True

34weight\_decay:float=0.0

36weight\_decay\_absolute:bool=False

39optimized\_adam\_update:bool=True

42parameters:any

45learning\_rate:float=0.01

47betas:Tuple[float,float]=(0.9,0.999)

49eps:float=1e-08

52momentum:float=0.5

54amsgrad:bool=False

57warmup:int=2\_000

59total\_steps:int=int(1e10)

62degenerate\_to\_sgd:bool=True

65rectify:bool=True

68d\_model:int

70def\_\_init\_\_(self):71super().\_\_init\_\_(\_primary='optimizer')727374meta\_config(OptimizerConfigs.parameters)

77@option(OptimizerConfigs.weight\_decay\_obj,'L2')78def\_weight\_decay(c:OptimizerConfigs):79returnWeightDecay(c.weight\_decay,c.weight\_decouple,c.weight\_decay\_absolute)808182@option(OptimizerConfigs.optimizer,'SGD')83def\_sgd\_optimizer(c:OptimizerConfigs):84returntorch.optim.SGD(c.parameters,c.learning\_rate,c.momentum,85weight\_decay=c.weight\_decay)868788@option(OptimizerConfigs.optimizer,'Adam')89def\_adam\_optimizer(c:OptimizerConfigs):90ifc.amsgrad:91fromlabml\_nn.optimizers.amsgradimportAMSGrad92returnAMSGrad(c.parameters,93lr=c.learning\_rate,betas=c.betas,eps=c.eps,94optimized\_update=c.optimized\_adam\_update,95weight\_decay=c.weight\_decay\_obj,amsgrad=c.amsgrad)96else:97fromlabml\_nn.optimizers.adamimportAdam98returnAdam(c.parameters,99lr=c.learning\_rate,betas=c.betas,eps=c.eps,100optimized\_update=c.optimized\_adam\_update,101weight\_decay=c.weight\_decay\_obj)102103104@option(OptimizerConfigs.optimizer,'AdamW')105def\_adam\_warmup\_optimizer(c:OptimizerConfigs):106fromlabml\_nn.optimizers.adam\_warmupimportAdamWarmup107returnAdamWarmup(c.parameters,108lr=c.learning\_rate,betas=c.betas,eps=c.eps,109weight\_decay=c.weight\_decay\_obj,amsgrad=c.amsgrad,warmup=c.warmup)110111112@option(OptimizerConfigs.optimizer,'RAdam')113def\_radam\_optimizer(c:OptimizerConfigs):114fromlabml\_nn.optimizers.radamimportRAdam115returnRAdam(c.parameters,116lr=c.learning\_rate,betas=c.betas,eps=c.eps,117weight\_decay=c.weight\_decay\_obj,amsgrad=c.amsgrad,118degenerated\_to\_sgd=c.degenerate\_to\_sgd)119120121@option(OptimizerConfigs.optimizer,'AdaBelief')122def\_ada\_belief\_optimizer(c:OptimizerConfigs):123fromlabml\_nn.optimizers.ada\_beliefimportAdaBelief124returnAdaBelief(c.parameters,125lr=c.learning\_rate,betas=c.betas,eps=c.eps,126weight\_decay=c.weight\_decay\_obj,amsgrad=c.amsgrad,127degenerate\_to\_sgd=c.degenerate\_to\_sgd,128rectify=c.rectify)129130131@option(OptimizerConfigs.optimizer,'Noam')132def\_noam\_optimizer(c:OptimizerConfigs):133fromlabml\_nn.optimizers.noamimportNoam134returnNoam(c.parameters,135lr=c.learning\_rate,betas=c.betas,eps=c.eps,136weight\_decay=c.weight\_decay\_obj,amsgrad=c.amsgrad,warmup=c.warmup,137d\_model=c.d\_model)138139140@option(OptimizerConfigs.optimizer,'AdamWarmupCosineDecay')141def\_noam\_optimizer(c:OptimizerConfigs):142fromlabml\_nn.optimizers.adam\_warmup\_cosine\_decayimportAdamWarmupCosineDecay143returnAdamWarmupCosineDecay(c.parameters,144lr=c.learning\_rate,betas=c.betas,eps=c.eps,145weight\_decay=c.weight\_decay\_obj,amsgrad=c.amsgrad,146warmup=c.warmup,total\_steps=c.total\_steps)