<jupyter_start><jupyter_text>Optimisation Methods for Deep Learning> In this article we will look at methods to improve gradient decent optimisation for training neural networks beyond SGD including momentum, RMSProp and Adam.- toc: true - comments: true- image: images/optimisation.jpeg- categories: [deep-learning-theory]<jupyter_code>#hide
!pip install -Uqq fastbook
import fastbook
fastbook.setup_book()
from fastbook import *<jupyter_output><empty_output><jupyter_text>IntroductionIn this article we will look at methods to improve gradient decent optimisation for training neural networks beyond SGD. These include momentum, RMSProp and Adam. We will also look at the fastai library system of callbacks which make changes to the training loop easier.This article is based on content from the [fastai deep learning course, chapter 16](https://github.com/fastai/fastbook/blob/master/16_accel_sgd.ipynb).  Basic SGDWe will first define a baseline using basic SGD to compare how further enhancements improve results. We will use the fastai curated *imagenette* dataset here.<jupyter_code>#hide_input
def get_data(url, presize, resize):
    path = untar_data(url)
    return DataBlock(
        blocks=(ImageBlock, CategoryBlock), get_items=get_image_files, 
        splitter=GrandparentSplitter(valid_name='val'),
        get_y=parent_label, item_tfms=Resize(presize),
        batch_tfms=[*aug_transforms(min_scale=0.5, size=resize),
                    Normalize.from_stats(*imagenet_stats)],
    ).dataloaders(path, bs=128)
dls = get_data(URLs.IMAGENETTE_160, 160, 128)<jupyter_output>A new version of this dataset is available, downloading...<jupyter_text>We will also create an untrained ResNet-34 architecture for our model which we will train from scratch.<jupyter_code>def get_learner(**kwargs):
    return cnn_learner(dls, resnet34, pretrained=False,
                    metrics=accuracy, **kwargs).to_fp16()

learn = get_learner()
learn.fit_one_cycle(3, 0.003)<jupyter_output><empty_output><jupyter_text>That was with all the default settings used by fastai. Lets explicitly use just basic SGD.<jupyter_code>learn = get_learner(opt_func=SGD)
learn.lr_find()<jupyter_output><empty_output><jupyter_text>So we will need to use a higher learning rate than we normally use. We will also need to explictly turn momentum off, as we are here trying to illustrate just using basic SGD.<jupyter_code>learn.fit_one_cycle(3, 0.03, moms=(0,0,0))<jupyter_output><empty_output><jupyter_text>Defining a generic optimiserThe fastai library provides a flexible approach to optimisers that makes it easier to add custom changes using *optimiser callbacks*. A key part of this is the *Optimiser* class which includes these two methods.<jupyter_code>def zero_grad(self):
    for p,*_ in self.all_params():
        p.grad.detach_()
        p.grad.zero_()

def step(self):
    for p,pg,state,hyper in self.all_params():
        for cb in self.cbs:
            state = _update(state, cb(p, **{**state, **hyper}))
        self.state[p] = state<jupyter_output><empty_output><jupyter_text>zero_grad is handy for clearing all the gradients. Note the step method loops through other potential callbacks which is how different aspects of optimisation old and new are done. Even basic SGD is one of these callbacks.<jupyter_code>def sgd_cb(p, lr, **kwargs): p.data.add_(-lr, p.grad.data)<jupyter_output><empty_output><jupyter_text>We can add this as a callback like this.<jupyter_code>opt_func = partial(Optimizer, cbs=[sgd_cb])<jupyter_output><empty_output><jupyter_text>Let's now train with this.<jupyter_code>learn = get_learner(opt_func=opt_func)
learn.fit(3, 0.03)<jupyter_output><empty_output><jupyter_text>MomentumSo the idea of Momentum is we want to go faster in the direction we are going with gradient decent to get there sooner. We could for example use a moving average.weight.avg = beta * weight.avg + (1-beta) * weight.gradnew_weight = weight - lr * weight.avg*beta* helps control how much momentum to use, so if its zero there is no momentum and we have just basic SGD. But if closer to 1 then the main direction is the average of the previous steps.<jupyter_code>#hide_input
#id img_momentum
#caption An example of momentum
#alt Graph showing an example of momentum
x = np.linspace(-4, 4, 100)
y = 1 - (x/3) ** 2
x1 = x + np.random.randn(100) * 0.1
y1 = y + np.random.randn(100) * 0.1
plt.scatter(x1,y1)
idx = x1.argsort()
beta,avg,res = 0.7,0,[]
for i in idx:
    avg = beta * avg + (1-beta) * y1[i]
    res.append(avg/(1-beta**(i+1)))
plt.plot(x1[idx],np.array(res), color='red');<jupyter_output><empty_output><jupyter_text>High beta can help us get over small 'bumps' in the loss landscape and keep going faster in the general direction of progress followed so far, but if too high can cause us to overshoot completly.<jupyter_code>#hide_input
#id img_betas
#caption Momentum with different beta values
#alt Graph showing how the beta value influences momentum
x = np.linspace(-4, 4, 100)
y = 1 - (x/3) ** 2
x1 = x + np.random.randn(100) * 0.1
y1 = y + np.random.randn(100) * 0.1
_,axs = plt.subplots(2,2, figsize=(12,8))
betas = [0.5,0.7,0.9,0.99]
idx = x1.argsort()
for beta,ax in zip(betas, axs.flatten()):
    ax.scatter(x1,y1)
    avg,res = 0,[]
    for i in idx:
        avg = beta * avg + (1-beta) * y1[i]
        res.append(avg)#/(1-beta**(i+1)))
    ax.plot(x1[idx],np.array(res), color='red');
    ax.set_title(f'beta={beta}')<jupyter_output><empty_output><jupyter_text>Beta too high means we really miss important changes in direction.fit_one_cycle starts with a high beta of 0.95, going down to 0.85 then back up to 0.95.Let's add momentum, by keeping track of the moving average gradient, which we can do with another callback.<jupyter_code>def average_grad(p, mom, grad_avg=None, **kwargs):
    if grad_avg is None: grad_avg = torch.zeros_like(p.grad.data)
    return {'grad_avg': grad_avg*mom + p.grad.data}

def momentum_step(p, lr, grad_avg, **kwargs): p.data.add_(-lr, grad_avg)

opt_func = partial(Optimizer, cbs=[average_grad,momentum_step], mom=0.9)<jupyter_output><empty_output><jupyter_text>Note Learner will automatically schedule the momentum and learning rate mom and lr, so fit_one_cycle will even work with our custom Optimiser.<jupyter_code>learn = get_learner(opt_func=opt_func)
learn.fit_one_cycle(3, 0.03)
learn.recorder.plot_sched()<jupyter_output><empty_output><jupyter_text>RMSPropRMSProp uses an adaptive learning rate, each parameter gets its own learning rate controlled by a global learning rate. The individual learning rate can be determined by looking at the gradients, for example if the gradients are close to zero for a while it might need a higher learning rate, and vice versa if the gradients are too high or unstable.We can use a moving average to get the general direction, specifically a moving average of the gradients squared.w.square_avg = alpha * w.square_avg + (1-alpha) * (w.grad ** 2)new_w = w - lr * w.grad / math.sqrt(w.square_avg + eps)The eps (epsilon) is added for numerical stability (usually set at 1e-8), and the default value for alpha is usually 0.99.<jupyter_code>def average_sqr_grad(p, sqr_mom, sqr_avg=None, **kwargs):
    if sqr_avg is None: sqr_avg = torch.zeros_like(p.grad.data)
    return {'sqr_avg': sqr_mom*sqr_avg + (1-sqr_mom)*p.grad.data**2}

def rms_prop_step(p, lr, sqr_avg, eps, grad_avg=None, **kwargs):
    denom = sqr_avg.sqrt().add_(eps)
    p.data.addcdiv_(-lr, p.grad, denom)

opt_func = partial(Optimizer, cbs=[average_sqr_grad,rms_prop_step],
                   sqr_mom=0.99, eps=1e-7)

learn = get_learner(opt_func=opt_func)
learn.fit_one_cycle(3, 0.003)<jupyter_output><empty_output><jupyter_text>AdamAdam combines SGD, momentum and RMSProp together. One difference is Adam uses an unbiased moving average.w.avg = beta * w.avg + (1-beta) * w.gradunbias_avg = w.avg / (1 - (beta**(i+1)))With all the steps combined we have:w.avg = beta1 * w.avg + (1-beta1) * w.gradunbias_avg = w.avg / (1 - (beta1**(i+1)))w.sqr_avg = beta2 * w.sqr_avg + (1-beta2) * (w.grad ** 2)new_w = w - lr * unbias_avg / sqrt(w.sqr_avg + eps)Adam is the default optimiser in fastai.  Decoupled Weight DecayWhen using Adam, we need to use a different kind of weight decay. Recall basic weight decay.new_weight = weight - lr*weight.grad - lr*wd*weightAnd alternative formulation is:weight.grad += wd*weightWith SGD these are the same, but not for Adam. So we need to use decoupled weight decay when using Adam.  Fastai CallbacksFastai callbacks allow you to add custom behaviour to the training loop at any point.![](https://github.com/pranath/blog/raw/master/images/fastai_loop.png "The fastai layered API")This has enabled easier adding of many new custom changes such as the below examples.![](https://github.com/pranath/blog/raw/master/images/fastai_callbacks.png "The fastai layered API")  Creating a CallbackLet's try defining a model reset callback.<jupyter_code>class ModelResetter(Callback):
    def begin_train(self):    self.model.reset()
    def begin_validate(self): self.model.reset()<jupyter_output><empty_output><jupyter_text>Here is another example RNN regulariser callback.<jupyter_code>class RNNRegularizer(Callback):
    def __init__(self, alpha=0., beta=0.): self.alpha,self.beta = alpha,beta

    def after_pred(self):
        self.raw_out,self.out = self.pred[1],self.pred[2]
        self.learn.pred = self.pred[0]

    def after_loss(self):
        if not self.training: return
        if self.alpha != 0.:
            self.learn.loss += self.alpha * self.out[-1].float().pow(2).mean()
        if self.beta != 0.:
            h = self.raw_out[-1]
            if len(h)>1:
                self.learn.loss += self.beta * (h[:,1:] - h[:,:-1]
                                               ).float().pow(2).mean()<jupyter_output><empty_output><jupyter_text>Inside the callback you can access global variables and objects such as self.model.  Callback Ordering and ExceptionsCallbacks can also interrupt any part of the training loop by using a system of exceptions, for example to skip a batch or stop training completely.This callback will stop training any time the loss becomes infinate.<jupyter_code>class TerminateOnNaNCallback(Callback):
    run_before=Recorder
    def after_batch(self):
        if torch.isinf(self.loss) or torch.isnan(self.loss):
            raise CancelFitException<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10523786641037962, 'fraction_numerical': 0.017011052378664102, 'mean_word_length': 4.529224229543039, 'pattern_counts': {'":': 0, '<': 68, '<?xml version=': 0, '>': 70, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '17', 'dir': 'jupyter-structured-clean-dedup', 'id': '616391', 'license': 'Apache-2.0', 'n_tokens_mistral': 3522, 'n_tokens_neox': 3371, 'n_words': 1067, 'path': '_notebooks/2021-06-13-optimisation-methods-for-deep-learning.ipynb', 'repo_name': 'pranath/blog', 'size': '10405'}