EigenPro-tensorflow/optimizers.py at master · EigenPro/EigenPro-tensorflow · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from keras import backend as K
from keras.optimizers import Optimizer

from backend_extra import scatter_update


class PSGD(Optimizer):
    """Primal Stochastic gradient descent optimizer.

    Arguments:
        pred_t: tensor. Prediction result.
        index_t: tensor. Mini-batch indices for primal updates.
        eta: float >= 0. Step size.
        eigenpro_f: Map grad tensor to EigenPro component.
    """

    def __init__(self, pred_t, index_t, eta=0.01, eigenpro_f=None, **kwargs):
        super(PSGD, self).__init__(**kwargs)
        self.eta = K.variable(eta, name='eta')
        self.pred_t = pred_t
        self.index_t = index_t
        self.eigenpro_f = eigenpro_f

    def get_updates(self, loss, params):
        self.updates = []
        grads = self.get_gradients(loss, [self.pred_t])

        eta = self.eta
        index = self.index_t
        eigenpro_f = self.eigenpro_f

        shapes = [K.get_variable_shape(p) for p in params]
        for p, g in zip(params, grads):
            update_p = K.gather(p, index) - eta * g
            new_p = scatter_update(p, index, update_p)

            if eigenpro_f:
                new_p = new_p + eta * eigenpro_f(g)

            self.updates.append(K.update(p, new_p))
        return self.updates

    def get_config(self):
        config = {'eta': float(K.get_value(self.eta))}
        base_config = super(PSGD, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class SGD(Optimizer):
    """Stochastic gradient descent optimizer.

    Arguments:
        eta: float >= 0. Step size.
        eigenpro_f: Map grad tensor to EigenPro component.
    """

    def __init__(self, eta=0.01, eigenpro_f=None, **kwargs):
        super(SGD, self).__init__(**kwargs)
        self.eta = K.variable(eta, name='eta')
        self.eigenpro_f = eigenpro_f

    def get_updates(self, loss, params):
        self.updates = []
        grads = self.get_gradients(loss, params)

        eta = self.eta
        eigenpro_f = self.eigenpro_f

        shapes = [K.get_variable_shape(p) for p in params]
        for p, g in zip(params, grads):
            new_p = p - eta * g

            if eigenpro_f:
                new_p = new_p + eta * eigenpro_f(g)

            self.updates.append(K.update(p, new_p))
        return self.updates

    def get_config(self):
        config = {'eta': float(K.get_value(self.eta))}
        base_config = super(SGD, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))