mirror of
https://github.com/linyiLYi/street-fighter-ai.git
synced 2025-04-04 23:20:43 +00:00
94 lines
4.2 KiB
Python
94 lines
4.2 KiB
Python
import torch
|
|
from torch.optim import Optimizer
|
|
|
|
class RMSpropTF(Optimizer):
|
|
def __init__(self, params, lr=1e-2, alpha=0.9, eps=1e-10,
|
|
weight_decay=0, momentum=0., centered=False,
|
|
decoupled_decay=False, lr_in_momentum=True
|
|
):
|
|
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,
|
|
decoupled_decay=decoupled_decay,
|
|
lr_in_momentum=lr_in_momentum
|
|
)
|
|
super(RMSpropTF, self).__init__(params, defaults)
|
|
|
|
def __setstate__(self, state):
|
|
super(RMSpropTF, self).__setstate__(state)
|
|
for group in self.param_groups:
|
|
group.setdefault('momentum', 0)
|
|
group.setdefault('centered', False)
|
|
|
|
@torch.no_grad()
|
|
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:
|
|
with torch.enable_grad():
|
|
loss = closure()
|
|
for group in self.param_groups:
|
|
for p in group['params']:
|
|
if p.grad is None:
|
|
continue
|
|
grad = p.grad
|
|
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.ones_like(p) # PyTorch inits to zero
|
|
if group['momentum'] > 0:
|
|
state['momentum_buffer'] = torch.zeros_like(p)
|
|
if group['centered']:
|
|
state['grad_avg'] = torch.zeros_like(p)
|
|
square_avg = state['square_avg']
|
|
one_minus_alpha = 1. - group['alpha']
|
|
state['step'] += 1
|
|
if group['weight_decay'] != 0:
|
|
if group['decoupled_decay']:
|
|
p.mul_(1. - group['lr'] * group['weight_decay'])
|
|
else:
|
|
grad = grad.add(p, alpha=group['weight_decay'])
|
|
|
|
# Tensorflow order of ops for updating squared avg
|
|
square_avg.add_(grad.pow(2) - square_avg, alpha=one_minus_alpha)
|
|
# square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha) # PyTorch original
|
|
if group['centered']:
|
|
grad_avg = state['grad_avg']
|
|
grad_avg.add_(grad - grad_avg, alpha=one_minus_alpha)
|
|
avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).add(group['eps']).sqrt_() # eps in sqrt
|
|
# grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha) #
|
|
# PyTorch original
|
|
else:
|
|
avg = square_avg.add(group['eps']).sqrt_() # eps moved in sqrt
|
|
if group['momentum'] > 0:
|
|
buf = state['momentum_buffer']
|
|
# Tensorflow accumulates the LR scaling in the momentum buffer
|
|
if group['lr_in_momentum']:
|
|
buf.mul_(group['momentum']).addcdiv_(grad, avg, value=group['lr'])
|
|
p.add_(-buf)
|
|
else:
|
|
# PyTorch scales the param update by LR
|
|
buf.mul_(group['momentum']).addcdiv_(grad, avg)
|
|
p.add_(buf, alpha=-group['lr'])
|
|
else:
|
|
p.addcdiv_(grad, avg, value=-group['lr'])
|
|
return loss
|
|
|