PyTorch Implementation of Gravity Optimizer

username: d4riush
cooked: <p>Hi, a few months ago my friend and I wrote a <a href="https://arxiv.org/abs/2101.09192" rel="noopener nofollow ugc">paper</a> about a new optimizer called <strong>Gravity</strong>. Initially, we only wrote Keras implementation of that. But recently I fall in love with PyTorch and added PyTorch implementation of Gravity to its <a href="https://github.com/dariush-bahrami/gravity.optimizer" rel="noopener nofollow ugc">repository</a>. The algorithm of Gravity is shown below.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a2daf20ccaf5a90facac364b24f802382bbe383.png" data-download-href="https://discuss.pytorch.org/uploads/default/4a2daf20ccaf5a90facac364b24f802382bbe383" title="Gravity Optimizer - Algorithm - Readme version"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a2daf20ccaf5a90facac364b24f802382bbe383_2_509x500.png" alt="Gravity Optimizer - Algorithm - Readme version" data-base62-sha1="aAdd3i4TdjXmd1ZXGb3WV3plXML" width="509" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a2daf20ccaf5a90facac364b24f802382bbe383_2_509x500.png, https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a2daf20ccaf5a90facac364b24f802382bbe383.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a2daf20ccaf5a90facac364b24f802382bbe383.png 2x" data-dominant-color="EFEFEF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Gravity Optimizer - Algorithm - Readme version</span><span class="informations">750×736 66.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
and the PyTorch implementation of this algorithm is as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.optim import Optimizer

class Gravity(Optimizer):
    """PyTorch implementation of Gravity optimizer.

    Gravity optimizer uses a kinematic approach to minimize the cost function. More
    detail can be found in corresponding paper at https://arxiv.org/abs/2101.09192.

    Args:
        params (iterable): iterable of parameters to optimize or dicts defining
            parameter groups
        lr (float, optional): learning rate (default: 0.1)
        alpha (float, optional): alpha controls the V initialization (defaults to 0.01)
        beta (float, optional): beta will be used to compute running
            average of V (defaults to 0.9)
    """

    def __init__(self, params, lr=0.1, alpha=0.01, beta=0.9):
        if not 0.0 &lt;= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 &lt;= alpha:
            raise ValueError("Invalid alpha value: {}".format(alpha))
        if not 0.0 &lt;= beta &lt; 1.0:
            raise ValueError("Invalid beta value: {}".format(beta))
        defaults = dict(lr=lr, alpha=alpha, beta=beta, t=0)
        super(Gravity, self).__init__(params, defaults)
        self.__init_v()

    def __init_v(self):
        for group in self.param_groups:
            lr = group["lr"]
            alpha = group["alpha"]
            std = alpha / lr
            for p in group["params"]:
                state = self.state[p]
                state["v"] = torch.empty_like(p).normal_(mean=0.0, std=std)

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.

        Args:
            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:
            lr = group["lr"]
            alpha = group["alpha"]
            beta = group["beta"]
            group["t"] += 1
            t = group["t"]
            beta_hat = (beta * t + 1) / (t + 2)
            for p in group["params"]:
                if p.grad is not None:
                    if p.grad.is_sparse:
                        raise RuntimeError("Gravity does not support sparse gradients")
                    state = self.state[p]
                    m = 1 / torch.abs(p.grad).max()
                    zeta = p.grad / (1 + (p.grad / m) ** 2)
                    state["v"] = beta_hat * state["v"] + (1 - beta_hat) * zeta
                    p -= lr * state["v"]
        return loss
</code></pre>
<p>I am not fluent in PyTorch yet so if you find any mistake in my code I will appreciate it if you let me know. Also, I will be happy if you give me some feedback on Gravity’s performance in your applications.</p>