Gradient descent is one of the most popular optimization methods. At present, gradient descent optimizers, combined with the loss function, are the key to machine learning, especially deep learning. In this section, we are going to understand:

• how to use optimizers in BrainPy?

• how to customize your own optimizer?

import brainpy as bp
import brainpy.math as bm

bp.math.set_platform('cpu')

import matplotlib.pyplot as plt


## Optimizers in BrainPy#

The basic optimizer class in BrainPy is brainpy.optimizers.Optimizer, which inludes the following optimizers:

• SGD

• Momentum

• Nesterov momentum

• RMSProp

All supported optimizers can be inspected through the brainpy.math.optimizers APIs.

Generally, an optimizer initialization receives the learning rate lr, the trainable variables train_vars, and other hyperparameters for the specific optimizer.

• lr can be a float, or an instance of brainpy.optim.Scheduler.

• train_vars should be a dict of Variable.

Here we launch a SGD optimizer.

a = bm.Variable(bm.ones((5, 4)))
b = bm.Variable(bm.zeros((3, 3)))

op = bp.optim.SGD(lr=0.001, train_vars={'a': a, 'b': b})


When you try to update the parameters, you must provide the corresponding gradients for each parameter in the update() method.

op.update({'a': bm.random.random(a.shape), 'b': bm.random.random(b.shape)})

print('a:', a)
print('b:', b)

a: Variable([[0.9998637 , 0.99998003, 0.99998164, 0.99945515],
[0.99987453, 0.9994551 , 0.9992173 , 0.99931335],
[0.99995905, 0.999395  , 0.9999578 , 0.9992222 ],
[0.9990663 , 0.9991484 , 0.99950355, 0.9991641 ],
[0.999546  , 0.99927944, 0.9995042 , 0.9993433 ]],            dtype=float32)
b: Variable([[-2.2797585e-06, -8.2317321e-04, -3.6020565e-04],
[-4.7648646e-04, -4.1223815e-04, -4.2962815e-04],
[-6.4605832e-05, -9.0033154e-04, -2.7676427e-04]],            dtype=float32)


You can process the gradients before applying them. For example, we clip the graidents by the maximum L2-norm.

grads_pre = {'a': bm.random.random(a.shape), 'b': bm.random.random(b.shape)}


{'a': JaxArray([[0.34498155, 0.5973849 , 0.3798045 , 0.9029752 ],
[0.48180568, 0.85238564, 0.9597529 , 0.19037902],
[0.09431374, 0.03752387, 0.05545044, 0.18625283],
[0.4265157 , 0.13127708, 0.44978166, 0.6449729 ],
[0.5855427 , 0.02250564, 0.9523196 , 0.317971  ]],            dtype=float32),
'b': JaxArray([[0.4190004 , 0.7033491 , 0.13393831],
[0.11366987, 0.33574808, 0.37153232],
[0.39718974, 0.25615263, 0.08950627]], dtype=float32)}

grads_post = bm.clip_by_norm(grads_pre, 1.)


{'a': JaxArray([[0.14619358, 0.2531551 , 0.16095057, 0.38265574],
[0.20417583, 0.3612173 , 0.40671656, 0.08067733],
[0.03996754, 0.01590157, 0.02349835, 0.07892876],
[0.18074548, 0.05563157, 0.19060494, 0.27332157],
[0.24813668, 0.00953726, 0.40356654, 0.13474725]],            dtype=float32),
'b': JaxArray([[0.38518783, 0.6465901 , 0.12312973],
[0.10449692, 0.30865383, 0.34155035],
[0.36513725, 0.23548158, 0.08228327]], dtype=float32)}

op.update(grads_post)

print('a:', a)
print('b:', b)

a: Variable([[0.9997175 , 0.9997269 , 0.9998207 , 0.9990725 ],
[0.9996703 , 0.9990939 , 0.9988105 , 0.99923265],
[0.99991906, 0.9993791 , 0.9999343 , 0.9991433 ],
[0.9988856 , 0.9990928 , 0.99931294, 0.99889076],
[0.99929786, 0.9992699 , 0.9991006 , 0.9992085 ]],            dtype=float32)
b: Variable([[-0.00038747, -0.00146976, -0.00048334],
[-0.00058098, -0.00072089, -0.00077118],
[-0.00042974, -0.00113581, -0.00035905]], dtype=float32)


Note

Optimizer usually has their own dynamically changed variables. If you JIT a function whose logic contains optimizer update, your dyn_vars in bm.jit() should include variables in Optimzier.vars().

op.vars()  # SGD optimzier only has an iterable step variable to record the training step

{'Constant0.step': Variable([2], dtype=int32)}

bp.optim.Momentum(lr=0.001, train_vars={'a': a, 'b': b}).vars()  # Momentum has velocity variables

{'Constant1.step': Variable([0], dtype=int32),
'Momentum0.a_v': Variable([[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]], dtype=float32),
'Momentum0.b_v': Variable([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]], dtype=float32)}

bp.optim.Adam(lr=0.001, train_vars={'a': a, 'b': b}).vars()  # Adam has more variables

{'Constant2.step': Variable([0], dtype=int32),
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]], dtype=float32),
[0., 0., 0.],
[0., 0., 0.]], dtype=float32),
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]], dtype=float32),
[0., 0., 0.],
[0., 0., 0.]], dtype=float32)}


## Creating A Self-Customized Optimizer#

To create your own optimization algorithm, simply inherit from bm.optimizers.Optimizer class and override the following methods:

• __init__(): init function that receives the learning rate (lr) and trainable variables (train_vars). Do not forget to register your dynamical changed variables into implicit_vars.

• update(grads): update function that computes the updated parameters.

The general structure is shown below:

class CustomizeOp(bp.optim.Optimizer):
def __init__(self, lr, train_vars, *params, **other_params):
super(CustomizeOp, self).__init__(lr, train_vars)

pass


## Schedulers#

Scheduler seeks to adjust the learning rate during training through reducing the learning rate according to a pre-defined schedule. Common learning rate schedules include time-based decay, step decay and exponential decay.

Here we set up an exponential decay scheduler, in which the learning rate will decay exponentially along the training step.

sc = bp.optim.ExponentialDecay(lr=0.1, decay_steps=2, decay_rate=0.99)

def show(steps, rates):
plt.plot(steps, rates)
plt.xlabel('Train Step')
plt.ylabel('Learning Rate')
plt.show()

steps = bm.arange(1000)
rates = sc(steps)

show(steps, rates)


After Optimizer initialization, the learning rate self.lr will always be an instance of bm.optimizers.Scheduler. A scalar float learning rate initialization will result in a Constant scheduler.

op.lr

Constant(0.001)


One can get the current learning rate value by calling Scheduler.__call__(i=None).

• If i is not provided, the learning rate value will be evaluated at the built-in training step.

• Otherwise, the learning rate value will be evaluated at the given step i.

op.lr()

0.001


In BrainPy, several commonly used learning rate schedulers are used:

• Constant

• ExponentialDecay

• InverseTimeDecay

• PolynomialDecay

• PiecewiseConstant

For more details, please see the brainpy.math.optimizers APIs.

# InverseTimeDecay scheduler

rates = bp.optim.InverseTimeDecay(lr=0.01, decay_steps=10, decay_rate=0.999)(steps)
show(steps, rates)

# PolynomialDecay scheduler

rates = bp.optim.PolynomialDecay(lr=0.01, decay_steps=10, final_lr=0.0001)(steps)
show(steps, rates)


## Creating a Self-Customized Scheduler#

If users try to implement their own scheduler, simply inherit from bm.optimizers.Scheduler class and override the following methods:

• __init__(): the init function.

• __call__(i=None): the learning rate value evalution.

class CustomizeScheduler(bp.optim.Scheduler):
def __init__(self, lr, *params, **other_params):
super(CustomizeScheduler, self).__init__(lr)