656 lines
19 KiB
Python
656 lines
19 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torch.optim as optim
|
|
|
|
from scepter.modules.opt.optimizers.base_optimizer import BaseOptimize
|
|
from scepter.modules.opt.optimizers.registry import OPTIMIZERS
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
|
|
SUPPORT_TYPES = ('Adadelta', 'Adagrad', 'Adam', 'Adamax', 'AdamW', 'ASGD',
|
|
'LBFGS', 'RMSprop', 'Rprop', 'SGD', 'SparseAdam')
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class SGD(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 0.1,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'MOMENTUM': {
|
|
'value': 0,
|
|
'description': 'the momentum!'
|
|
},
|
|
'DAMPENING': {
|
|
'value': 0,
|
|
'description': 'the dampening!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
},
|
|
'NESTEROV': {
|
|
'value': False,
|
|
'description': 'the nesterov!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(SGD, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.LEARNING_RATE
|
|
self.momentum = cfg.get('MOMENTUM', 0)
|
|
self.dampening = cfg.get('DAMPENING', 0)
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
self.nesterov = cfg.get('NESTEROV', False)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.SGD(parameters,
|
|
lr=self.lr,
|
|
momentum=self.momentum,
|
|
dampening=self.dampening,
|
|
weight_decay=self.weight_decay,
|
|
nesterov=self.nesterov)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
SGD.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class Adadelta(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1.0,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'RHO': {
|
|
'value': 0,
|
|
'description': 'the rho!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-6,
|
|
'description': 'the eps!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Adadelta, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1.0)
|
|
self.rho = cfg.get('RHO', 0.0)
|
|
self.eps = cfg.get('EPS', 1e-6)
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.Adadelta(parameters,
|
|
lr=self.lr,
|
|
rho=self.rho,
|
|
eps=self.eps,
|
|
weight_decay=self.weight_decay)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
Adadelta.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class Adagrad(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1e-2,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'LEARNING_RATE_DECAY': {
|
|
'value': 0,
|
|
'description': 'the lr decay!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-10,
|
|
'description': 'the eps!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
},
|
|
'INITIAL_ACCUMULATOR_VALUE': {
|
|
'value': 0,
|
|
'description': 'the initial accumulator value!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Adagrad, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-2)
|
|
self.lr_decay = cfg.get('LEARNING_RATE_DECAY', 0.0)
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
self.initial_accumulator_value = cfg.get('INITIAL_ACCUMULATOR_VALUE',
|
|
0)
|
|
self.eps = cfg.get('EPS', 1e-10)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.Adagrad(
|
|
parameters,
|
|
lr=self.lr,
|
|
lr_decay=self.lr_decay,
|
|
weight_decay=self.weight_decay,
|
|
initial_accumulator_value=self.initial_accumulator_value,
|
|
eps=self.eps)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
Adagrad.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class Adam(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1e-2,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'BETAS': {
|
|
'value': [0.9, 0.999],
|
|
'description': 'the rho!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-6,
|
|
'description': 'the eps!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
},
|
|
'AMSGRAD': {
|
|
'value': False,
|
|
'description': 'the amsgrad!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Adam, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-2)
|
|
self.betas = cfg.get('BETAS', [0.9, 0.999])
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
self.amsgrad = cfg.get('AMSGRAD', False)
|
|
self.eps = cfg.get('EPS', 1e-10)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.Adam(parameters,
|
|
lr=self.lr,
|
|
betas=tuple(self.betas),
|
|
eps=self.eps,
|
|
weight_decay=self.weight_decay,
|
|
amsgrad=self.amsgrad)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
Adam.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class Adamax(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 2e-3,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'BETAS': {
|
|
'value': [0.9, 0.999],
|
|
'description': 'the rho!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-8,
|
|
'description': 'the eps!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Adamax, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 2e-3)
|
|
self.betas = cfg.get('BETAS', [0.9, 0.999])
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
self.eps = cfg.get('EPS', 1e-8)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.Adamax(parameters,
|
|
lr=self.lr,
|
|
betas=tuple(self.betas),
|
|
eps=self.eps,
|
|
weight_decay=self.weight_decay)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
Adamax.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class AdamW(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1e-3,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'BETAS': {
|
|
'value': [0.9, 0.999],
|
|
'description': 'the rho!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-8,
|
|
'description': 'the eps!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
},
|
|
'AMSGRAD': {
|
|
'value': False,
|
|
'description': 'the amsgrad!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(AdamW, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-3)
|
|
self.betas = cfg.get('BETAS', [0.9, 0.999])
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
self.eps = cfg.get('EPS', 1e-8)
|
|
self.amsgrad = cfg.get('AMSGRAD', False)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.AdamW(parameters,
|
|
lr=self.lr,
|
|
betas=tuple(self.betas),
|
|
eps=self.eps,
|
|
weight_decay=self.weight_decay,
|
|
amsgrad=self.amsgrad)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
AdamW.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class ASGD(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1e-2,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'LAMBD': {
|
|
'value': 1e-4,
|
|
'description': 'the rho!'
|
|
},
|
|
'ALPHA': {
|
|
'value': 0.75,
|
|
'description': 'the alpha!'
|
|
},
|
|
'T0': {
|
|
'value': 1e6,
|
|
'description': 'the t0!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(ASGD, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-2)
|
|
self.lambd = cfg.get('LAMBD', 1e-4)
|
|
self.alpha = cfg.get('ALPHA', 0.75)
|
|
self.t0 = cfg.get('T0', 1e6)
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', 0)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.ASGD(parameters,
|
|
lr=self.lr,
|
|
lambd=self.lambd,
|
|
alpha=self.alpha,
|
|
t0=self.t0,
|
|
weight_decay=self.weight_decay)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
ASGD.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class LBFGS(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1.0,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'MAX_ITER': {
|
|
'value': 20,
|
|
'description': 'the max iter!'
|
|
},
|
|
'MAX_EVAL': {
|
|
'value': None,
|
|
'description': 'the max eval!'
|
|
},
|
|
'TOLERANCE_GRAD': {
|
|
'value': 1e-7,
|
|
'description': 'the tolerance grad!'
|
|
},
|
|
'TOLERANCE_CHANGE': {
|
|
'value': 1e-9,
|
|
'description': 'the tolerance change!'
|
|
},
|
|
'HISTORY_SIZE': {
|
|
'value': 100,
|
|
'description': 'the history size!'
|
|
},
|
|
'LINE_SEARCH_FN': {
|
|
'value': None,
|
|
'description': 'the line search fn!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(LBFGS, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1)
|
|
self.max_iter = cfg.get('MAX_ITER', 20)
|
|
self.max_eval = cfg.get('MAX_EVAL', None)
|
|
self.tolerance_grad = cfg.get('TOLERANCE_GRAD', 1e-7)
|
|
self.tolerance_change = cfg.get('TOLERANCE_CHANGE', 1e-9)
|
|
self.history_size = cfg.get('HISTORY_SIZE', 100)
|
|
self.line_search_fn = cfg.get('LINE_SEARCH_FN', None)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.LBFGS(parameters,
|
|
lr=self.lr,
|
|
max_iter=self.max_iter,
|
|
max_eval=self.max_eval,
|
|
tolerance_grad=self.tolerance_grad,
|
|
tolerance_change=self.tolerance_change,
|
|
history_size=self.history_size,
|
|
line_search_fn=self.line_search_fn)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
LBFGS.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class RMSprop(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1e-2,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'ALPHA': {
|
|
'value': 0.99,
|
|
'description': 'the alpha!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-8,
|
|
'description': 'the eps!'
|
|
},
|
|
'WEIGHT_DECAY': {
|
|
'value': 0,
|
|
'description': 'the weight decay!'
|
|
},
|
|
'MOMENTUM': {
|
|
'value': 0,
|
|
'description': 'the momentum!'
|
|
},
|
|
'CENTERED': {
|
|
'value': False,
|
|
'description': 'the centered!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(RMSprop, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-2)
|
|
self.alpha = cfg.get('ALPHA', 0.99)
|
|
self.eps = cfg.get('EPS', 1e-8)
|
|
self.weight_decay = cfg.get('WEIGHT_DECAY', False)
|
|
self.momentum = cfg.get('MOMENTUM', 0)
|
|
self.centered = cfg.get('CENTERED', False)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.RMSprop(parameters,
|
|
lr=self.lr,
|
|
alpha=self.alpha,
|
|
eps=self.eps,
|
|
weight_decay=self.weight_decay,
|
|
momentum=self.momentum,
|
|
centered=self.centered)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
RMSprop.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class Rprop(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1.0,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'ETAS': {
|
|
'value': [0.5, 1.2],
|
|
'description': 'the etas!'
|
|
},
|
|
'STEP_SIZES': {
|
|
'value': [1e-6, 50],
|
|
'description': 'the step sizes!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Rprop, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-2)
|
|
self.etas = cfg.get('ETAS', [0.5, 1.2])
|
|
self.step_sizes = cfg.get('STEP_SIZES', [1e-6, 50])
|
|
|
|
def __call__(self, parameters):
|
|
return optim.Rprop(parameters,
|
|
lr=self.lr,
|
|
etas=tuple(self.etas),
|
|
step_sizes=tuple(self.step_sizes))
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
Rprop.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@OPTIMIZERS.register_class()
|
|
class SparseAdam(BaseOptimize):
|
|
para_dict = {
|
|
'LEARNING_RATE': {
|
|
'value': 1.0,
|
|
'description': 'the initial learning rate!'
|
|
},
|
|
'BETAS': {
|
|
'value': [0.9, 0.999],
|
|
'description': 'the betas!'
|
|
},
|
|
'EPS': {
|
|
'value': 1e-8,
|
|
'description': 'the eps!'
|
|
}
|
|
}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(SparseAdam, self).__init__(cfg, logger=logger)
|
|
self.lr = cfg.get('LEARNING_RATE', 1e-3)
|
|
self.betas = cfg.get('BETAS', [0.9, 0.999])
|
|
self.eps = cfg.get('EPS', 1e-8)
|
|
|
|
def __call__(self, parameters):
|
|
return optim.SparseAdam(parameters,
|
|
lr=self.lr,
|
|
betas=tuple(self.betas),
|
|
eps=self.eps)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('OPTIMIZER',
|
|
__class__.__name__,
|
|
SparseAdam.para_dict,
|
|
set_name=True)
|