KgeOptimizer now looks up the optimizer name directly in torch.optim.
Arbitrary options for that optimizer are now tunneled through with
unchecked options.
@@ -4,18 +4,13 @@ import torch.optim
class KgeOptimizer:
""" Wraps torch optimizers """
+ @staticmethod
def create(config, model):
""" Factory method for optimizer creation """
- if config.get('train.optimizer') == 'adagrad':
- return torch.optim.Adagrad(model.parameters(), config.get('train.lr'))
- elif config.get('train.optimizer') == 'adam':
- if config.get('model.sparse'):
- return torch.optim.SparseAdam(model.parameters(), config.get('train.lr'))
- else:
- return torch.optim.Adam(model.parameters(), config.get('train.lr'))
- elif config.get('train.optimizer') == 'sgd':
- return torch.optim.SGD(model.parameters(), config.get('train.lr'))
- else:
+ try:
+ optimizer = getattr(torch.optim, config.get('train.optimizer'))
+ return optimizer([p for p in model.parameters() if p.requires_grad], **config.get('train.optimizer_args'))
+ except:
# perhaps TODO: try class with specified name -> extensibility
raise ValueError('train.optimizer')
