Spaces:
Sleeping
Sleeping
| import torch | |
| from torch.optim import Optimizer | |
| class DM_RMSprop(Optimizer): | |
| """Implements the form of RMSProp used in DM 2015 Atari paper. | |
| Inspired by https://github.com/spragunr/deep_q_rl/blob/master/deep_q_rl/updates.py""" | |
| def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False): | |
| if not 0.0 <= lr: | |
| raise ValueError("Invalid learning rate: {}".format(lr)) | |
| if not 0.0 <= eps: | |
| raise ValueError("Invalid epsilon value: {}".format(eps)) | |
| if not 0.0 <= momentum: | |
| raise ValueError("Invalid momentum value: {}".format(momentum)) | |
| if not 0.0 <= weight_decay: | |
| raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) | |
| if not 0.0 <= alpha: | |
| raise ValueError("Invalid alpha value: {}".format(alpha)) | |
| defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, centered=centered, weight_decay=weight_decay) | |
| super(DM_RMSprop, self).__init__(params, defaults) | |
| def __setstate__(self, state): | |
| super(DM_RMSprop, self).__setstate__(state) | |
| for group in self.param_groups: | |
| group.setdefault('momentum', 0) | |
| group.setdefault('centered', False) | |
| def step(self, closure=None): | |
| """Performs a single optimization step. | |
| Arguments: | |
| closure (callable, optional): A closure that reevaluates the model | |
| and returns the loss. | |
| """ | |
| loss = None | |
| if closure is not None: | |
| loss = closure() | |
| for group in self.param_groups: | |
| momentum = group['momentum'] | |
| sq_momentum = group['alpha'] | |
| epsilon = group['eps'] | |
| for p in group['params']: | |
| if p.grad is None: | |
| continue | |
| grad = p.grad.data | |
| if grad.is_sparse: | |
| raise RuntimeError('RMSprop does not support sparse gradients') | |
| state = self.state[p] | |
| # State initialization | |
| if len(state) == 0: | |
| state['step'] = 0 | |
| state['square_avg'] = torch.zeros_like(p.data) | |
| if momentum > 0: | |
| state['momentum_buffer'] = torch.zeros_like(p.data) | |
| mom_buffer = state['momentum_buffer'] | |
| square_avg = state['square_avg'] | |
| state['step'] += 1 | |
| mom_buffer.mul_(momentum) | |
| mom_buffer.add_((1 - momentum) * grad) | |
| square_avg.mul_(sq_momentum).addcmul_(1 - sq_momentum, grad, grad) | |
| avg = (square_avg - mom_buffer**2 + epsilon).sqrt() | |
| p.data.addcdiv_(-group['lr'], grad, avg) | |
| return loss | |