Spaces:
Runtime error
Runtime error
File size: 673 Bytes
eea83e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | import torch.optim as optim
def get_optimizer(config, parameters):
if config.optim.optimizer == 'Adam':
return optim.Adam(parameters, lr=config.optim.lr, weight_decay=config.optim.weight_decay,
betas=(0.9, 0.999), amsgrad=config.optim.amsgrad, eps=config.optim.eps)
elif config.optim.optimizer == 'RMSProp':
return optim.RMSprop(parameters, lr=config.optim.lr, weight_decay=config.optim.weight_decay)
elif config.optim.optimizer == 'SGD':
return optim.SGD(parameters, lr=config.optim.lr, momentum=0.9)
else:
raise NotImplementedError('Optimizer {} not understood.'.format(config.optim.optimizer))
|