JittorMirror/python/jittor/optim.py

# ***************************************************************
# Copyright (c) 2020 Jittor. Authors:
#     Guowei Yang <471184555@qq.com>
#     Guoye Yang <498731903@qq.com>
#     Wenyang Zhou <576825820@qq.com>
#     Meng-Hao Guo <guomenghao1997@gmail.com>
#     Dun Liang <randonlang@gmail.com>.
#
# All Rights Reserved.
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
import jittor as jt
import numpy as np

class Optimizer(object):
    """ Basic class of Optimizer.
    Example:
    ```
    optimizer = nn.SGD(model.parameters(), lr)
    optimizer.step(loss)
    ```
    """
    def __init__(self, params, lr, param_sync_iter=10000):
        self.param_groups = []
        self.lr = lr
        self.param_sync_iter = param_sync_iter

        assert len(params) > 0, "Length of parameters should not be zero"
        if not isinstance(params[0], dict):
            params = [{'params': params}]
        for pg in params:
            assert isinstance(pg, dict)
            self.param_groups.append(pg)
        self.n_step = 0

    def pre_step(self, loss):
        """ something should be done before step, 
        such as calc gradients, mpi sync, and so on.
        Example:
        ```
        class MyOptimizer(Optimizer):
            def step(self, loss):
                self.post_step(loss)
                ...
        ```
        """
        # clean prev grads
        params = []
        params_has_grad = []
        for pg in self.param_groups:
            pg["grads"] = [None] * len(pg['params'])
            for p in pg['params']:
                params.append(p)
                if not p.is_stop_grad():
                    params_has_grad.append(p)
        
        # sync params, reduce computing graph size
        jt.sync(params)

        # get gradient
        grads = jt.grad(loss, params_has_grad)

        # sync grads and model if in mpi
        if jt.mpi:
            for g in grads:
                g.assign(g.mpi_all_reduce("mean"))
            if self.n_step % self.param_sync_iter == 0:
                for p in params:
                    p.assign(p.mpi_all_reduce("mean"))
        self.n_step += 1

        # set up grads in param_groups
        pid = 0
        for pg in self.param_groups:
            pg_grads = pg["grads"]
            for i, p in enumerate(pg['params']):
                if not p.is_stop_grad():
                    pg_grads[i] = grads[pid]
                    pid += 1
        
    def step(self, loss):
        self.pre_step(loss)
        for pg in self.param_groups:
            lr = pg.get("lr", self.lr)
            for p, g in zip(pg["params"], pg["grads"]):
                if p.is_stop_grad(): continue
                p -= g * lr
                # detach with the prev graph to reduce memory consumption
                p.detach_inplace()


class SGD(Optimizer):
    """ SGD Optimizer.
    Example:
    ```
    optimizer = nn.SGD(model.parameters(), lr, momentum=0.9)
    optimizer.step(loss)
    ```
    """
    def __init__(self, params, lr, momentum=0, weight_decay=0, dampening=0, nesterov=False):
        super().__init__(params, lr)
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.dampening = dampening
        self.nesterov = nesterov

        # initialize required arguments
        for pg in self.param_groups:
            values = pg["values"] = []
            for p in pg["params"]:
                values.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad())

    def step(self, loss):
        self.pre_step(loss)
        for pg in self.param_groups:
            # get arguments from each param_groups
            lr = pg.get("lr", self.lr)
            momentum = pg.get("momentum", self.momentum)
            weight_decay = pg.get("weight_decay", self.weight_decay)
            dampening = pg.get("dampening", self.dampening)
            nesterov = pg.get("nesterov", self.nesterov)

            # optimize main body
            for p, g, v in zip(pg["params"], pg["grads"], pg["values"]):
                if p.is_stop_grad(): continue
                dp = p * weight_decay + g
                v.assign(momentum * v + dp * (1 - dampening))
                if nesterov:
                    p -= (dp + momentum * v) * lr
                else:
                    p -= v * lr
                p.detach_inplace()

class Adam(Optimizer):
    """ Adam Optimizer.
    Example:
    ```
    optimizer = nn.Adam(model.parameters(), lr, eps=1e-8, betas=(0.9, 0.999))
    optimizer.step(loss)
    ```
    """
    def __init__(self, params, lr, eps=1e-8, betas=(0.9, 0.999), weight_decay=0):
        super().__init__(params, lr)
        self.eps = eps
        self.betas = betas
        # self.weight_decay = weight_decay
        assert weight_decay==0, "weight_decay is not supported yet"
        
        # initialize required arguments for each param_groups
        for pg in self.param_groups:
            values = pg["values"] = []
            m = pg["m"] = []
            for p in pg["params"]:
                values.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad())
                m.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad())

    def step(self, loss):
        self.pre_step(loss)
        n = float(self.n_step)
        for pg in self.param_groups:
            # get arguments from each param_groups
            lr = pg.get("lr", self.lr)
            eps = pg.get("eps", self.eps)
            b0, b1 = pg.get("betas", self.betas)
            for p, g, v, m in zip(pg["params"], pg["grads"], pg["values"], pg["m"]):
                if p.is_stop_grad(): continue
                m.assign(b0 * m + (1-b0) * g)
                v.assign(b1 * v + (1-b1) * g * g)
                step_size = lr * jt.sqrt(1-b1**n) / (1-b0 ** n)
                p -= m * step_size / (jt.sqrt(v) + eps)
                p.detach_inplace()