#!/usr/bin/env python3
# Copyright (c) 2018-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from torch.optim.optimizer import Optimizer, required
[docs]class RiemannianSGD(Optimizer):
"""Riemannian stochastic gradient descent.
Args:
rgrad (Function): Function to compute the Riemannian gradient
from the Euclidean gradient
retraction (Function): Function to update the retraction
of the Riemannian gradient
"""
def __init__(
self,
params,
lr=required,
rgrad=required,
expm=required,
):
defaults = {
'lr': lr,
'rgrad': rgrad,
'expm': expm,
}
super(RiemannianSGD, self).__init__(params, defaults)
[docs] def step(self, lr=None, **kwargs):
"""Performs a single optimization step.
Arguments:
lr (float, optional): learning rate for the current update.
"""
loss = None
for group in self.param_groups:
for p in group['params']:
lr = lr or group['lr']
rgrad = group['rgrad']
expm = group['expm']
if p.grad is None:
continue
d_p = p.grad.data
# make sure we have no duplicates in sparse tensor
if d_p.is_sparse:
d_p = d_p.coalesce()
d_p = rgrad(p.data, d_p)
d_p.mul_(-lr)
expm(p.data, d_p)
return loss