How to Easily Use Gradient Accumulation in Keras Models | by Raz Rotenberg | Towards Data Science

In another article, we covered what is gradient accumulation in deep learning and how it can solve issues when running neural networks with large batch sizes.
In this article, we will first see how you can easily use the generic gradient accumulation tool we implemented and used at Run:AI. Then, we will deep-dive into Keras optimizers and the way we have implemented such a generic tool.
The code is available on GitHub along with examples you can use right out of the box.
Adding gradient accumulation support to your Keras models is extremely simple. First of all, install the Run:AI Python library using the command:
pip install runai
Then, import the gradient accumulation package into your Python code:
import runai.ga
Now, you can choose one of two options. You can either wrap an existing optimizer with the generic gradient accumulation wrapper or create a gradient accumulation version of any of the built-in optimizers. The two options require specifying the number of steps you want to accumulate gradients over (passed as STEPS in the examples below).
Use the next line to wrap an existing optimizer (where optimizer is your optimizer):
optimizer = runai.ga.keras.optimizers.Optimizer(    optimizer, steps=STEPS)
Or, use the next line to create a gradient accumulation version of any of the built-in optimizers (“Adam” is used in the next example):
optimizer = runai.ga.keras.optimizers.Adam(steps=STEPS)
And that’s it! you have successfully added gradient accumulation support to your Keras model.
After seeing how easy it is to use, let’s now see what’s going on under the hood. First, we are going to examine the concept of optimizers in Keras and discuss their responsibility and their implementation. Then, we will deep-dive into how we implemented such a generic mechanism.
Optimizers in Keras are responsible for implementing the optimization algorithm — the mathematical formula responsible for minimizing the loss function. They receive all the model parameters — weights and biases — as inputs, compute their gradients and use them to generate updates for the model parameters. The updates for the model parameters are not the gradients themselves and are calculated using the gradients, as well as other parameters.
Every optimization algorithm has parameters. Some optimizers have similar parameters (e.g. “epsilon”, “beta_1”, “beta_2”, etc...), some may have unique ones (e.g. “amsgrad” in Adam), and they all support “learning rate” (and learning rate decay).
While building the model, Keras will call the optimizer and will pass two arguments. The first argument will be all the trainable variables — weights and biases — of the model. The second argument will be the loss value. Note that the trainable variables and the loss are tensors and not the actual values of the tensors.
The optimizer then takes these input tensors and adds the optimization algorithm to the model. It first calculates the gradients of the trainable variables with respect to the loss (by calling tf.gradients()), and then generates the tensors representing the mathematical formula.
Keras will then evaluate these tensors — that were generated by the optimizer — every step. By evaluating those tensors, the gradients will be calculated, and then the variable updates will be calculated and assigned to the model variables — the weights and biases.
Keras optimizers are Python classes, and they all inherit a base class called Optimizer (can be seen here). The base class implements a few methods and declares other (virtual) methods that must be implemented in subclasses. Let’s briefly examine the methods in Optimizer:
__init__(self, **kwargs): Initializes common members and configuration parameters.
get_updates(self, loss, params): A virtual method that is not implemented in the base class, and must be defined in subclasses, implementing the mathematical formula of the optimizer.
get_gradients(self, loss, params): A utility method to ease the gradient calculations of loss with respect to params (a wrapper of tf.gradients() technically).
set_weights(self, weights): Sets the values of the optimizer’s weights.
get_weights(self): Gets the values of the optimizer’s weights.
get_config(self): Gets the common configuration values.
We will be focusing on get_updates(). This method is only declared (virtual) in the base class; it is not implemented in it and must be implemented in all subclasses. Its implementation is the main difference between one optimizer to another.
As we said, Keras will call get_updates() as part of building the model. The method receives two arguments: loss and params, which both are tensors. and returns a list of tensors, which are the “Assign” ops — the tensors that actually assign the variable updates upon evaluation.
Let’s examine a simplified version of the implementation of SGD’s get_updates(). Note that this is a simplified version and not the actual code in Keras:
First, at line 2, the optimizer calculates the gradients of loss with respect to params by calling self.get_gradients().
Then, in lines 4–6, the optimizer iterates through all the trainable variables of the model, with their respective calculated gradients. For every parameter, it calculates (line 5) the new value of the variable (new_p) using the gradient and the learning rate (self.lr is a tensor initialized in __init__()). Then, it creates a tensor that will assign the variable’s new value (K.update(p, new_p)) and adds it to the list of such tensors (self.updates).
Finally, in line 8, the optimizer returns the list of the “Assign” tensors (self.updates), which will be evaluated by Keras in every step.
This method is called exactly once when the model is being built. It may take some time to digest, but be sure to understand that the arguments and the results are tensor objects and not the actual tensor values. Note that line 5 may be misleading and may seem like it calculates actual values, but this is not the case and is just syntactic sugar (TensorFlow’s operator overloading) on tensor objects.
After examining a simplified version, it’s time to figure out what happens in reality.
In addition to their plain algorithm, most of the optimizers support learning rate decay, which means they modify the learning rate throughout the training of the model, instead of having a constant value for it. In order to support modifying the learning rate over time, the concept of time should be defined. As expected, the learning rate is not modified as a function of the actual time that has passed since the beginning of the training phase, but as a function of the step number, and to support this, optimizers count steps (called iterations in their implementations).
Moreover, SGD in Keras supports momentums and the Nesterov momentum algorithm, which complicates things a bit more. Let’s take a look at the actual implementation of SGD’s get_updates() (can be seen on GitHub also):
Let’s go through the additions over the simplified version we examined before.
In line 3, the optimizer creates an “Assign” tensor to increase self.iterations by 1 in every step. self.iterations is the step counter (a tensor object) and is created in __init__() with an initial value of 0.
In lines 5–7 we can see the learning rate decay formula as a function of the step counter (self.iterations) and the decay value (an argument passed to the optimizer upon creation, and set to self.decay).
In line 11 it creates variables for the momentums. It declares another variable, for every trainable variable, with the same shape, and initializes it to 0 (the call to K.zeros(shape)).
Then, inside the iteration over the trainable variables, it calculates the new value for every momentum (line 14) and creates an “Assign” tensor to update the momentum with the new value (line 15).
Then, in lines 17–20, it calculates the new value for the parameter — depending on whether configured to apply Nesterov momentum or not — and applies the necessary constraints on the parameter (if there are such) in lines 23–24.
The rest of the lines are similar to the simplified version from before.
You can wander around optimizers.py and read the implementations of the different optimizers of Keras.
Now, after deep-diving into what exactly Keras optimizers are and how they are implemented, we are ready to discuss the different implementation alternatives for a gradient accumulation mechanism.
It is possible to rewrite any optimizer to support gradient-accumulation. Gradients should be accumulated over a few steps and only then should the optimizer use them for updating the model parameters. This is not optimal as gradient accumulation is a generic approach and should be optimizer-independent, and there are several flaws to this approach:
Every optimizer has a different formula for implementing a different optimization algorithm and therefore will require a different implementation for the respective gradient accumulation version.This is quite version-specific and code modification will be required every time the original implementation will be changed or a new optimizer will be added.It causes code duplication and is less elegant.
Every optimizer has a different formula for implementing a different optimization algorithm and therefore will require a different implementation for the respective gradient accumulation version.
This is quite version-specific and code modification will be required every time the original implementation will be changed or a new optimizer will be added.
It causes code duplication and is less elegant.
A preferable approach is to design the gradient accumulation model so that it can wrap any Keras optimizer regardless of its optimization algorithm.
By having a generic gradient accumulation mechanism, changes in the original optimizers will not require code updates.
In order to design and implement a generic gradient accumulation mechanism, there are some things that need to be taken into consideration.
Running the optimization algorithm on every mini-batch will not result in the same updates for the model parameters. In other words, we cannot just evaluate the optimization algorithm in every step — on every mini-batch. Otherwise, there was no need for gradient accumulation, and we could have just used a smaller batch size.
If we were to use the global batch, all the gradients would have been calculated using the same values of the model parameters — the weights and biases. When splitting the global batch into several mini-batches, evaluating the optimization algorithm every step will cause the model parameters to be updated after every mini-batch. This means that the gradients of all mini-batches will not be calculated using the same values of the weights and biases.
In addition, optimizers use various parameters as part of their formula, and those parameters are updated as part of the evaluation of the optimization algorithm. Updating those parameters every step — after every mini-batch — will result in changes in the state of the optimization algorithm between different mini-batches.
Our wrapper is a Python class that inherits Keras’s base Optimizer class. We receive the original optimizer as an argument upon creation (in __init__()), as well as the number of steps we want to accumulate gradients over.
We define all the methods exposed by optimizers (i.e. get_gradients(), set_weights(), get_weights(), etc...) and transparently call the original optimizer’s respective methods. The main logic lives — as expected — in get_updates().
Let’s start examining get_updates() (can be seen on GitHub as well), and deep-dive into the algorithm and implementation:
The first line (2) should seem familiar, where we calculate the gradients in the same way as other optimizers do. Note that grads will hold the value of the gradients of every mini-batch.
In line 5, we declare a step counter — called iterations — with an initial value of 0 (pretty similar to other optimizers). We use the step counter to tell if we are either at the first or the last step of the accumulation. To do so, we declare two tensors: first and last (lines 6–7).
first will be set to True every time we passed exactly self.steps steps. Technically, this is when iterations % self.steps will be equal to 0. For example, if we accumulate over five steps, this will be the case at the first step (indexed 0), the sixth step (indexed 5), the eleventh step (indexed 10), etc... In those cases, at the beginning of the step, we want to reset the accumulated gradients to 0 and start accumulating them once again.
last will be set to True every step we want to update the variables. Technically, this is when iterations % self.steps will be equal to self.steps — 1. Continuing the example from before, this will be the case at the fifth step (indexed 4), tenth step (indexed 9), etc...
In line 10, we declare variables to hold the values accumulated gradients between steps. We declare such a variable for every model parameter — for every trainable weight or bias — with the shape and type of the parameter, and initialize them with zeros.
Using those variables, in line 13 we declare tensors — agrads — to hold the values of accumulated gradients in every step. We use first to tell whether we should start accumulating the gradients from now on or use the gradients accumulated in the previous steps. If first is True — meaning we start accumulating from now on — we use the gradients of the current mini-batch alone. If first is False — meaning we should use the gradients accumulated over the past steps — we add the gradients of the current mini-batch to vagrads. This control flow (checking the value of first) is generated into the model using K.switch().
As a generic wrapper, we don’t implement any optimization algorithm. The original optimizer is responsible for that. As we covered, every optimizer implements its mathematical formula in get_updates(). There, the optimizer manages and uses all the needed parameters for the formula (e.g. step counter, learning rate, momentums, etc...). The optimizer stores the values of the parameters in dedicated variables, and every time a parameter needs to be updated, it assigns the new value to its dedicated variable.
The method get_updates() is called once, generating “Assign” tensors that will be evaluated in every step. Some of them are the updates of the model parameters, and the other ones are updates to the optimizer parameters.
As long as we are accumulating gradients, we don’t want any of these updates to happen. We don’t want the model parameters to be updated in order for all the mini-batches to start from the exact same point, in terms of having the same values for the weights and biases. We don’t want the optimizer parameters to be updated in order for the optimizer to advance in the pace as if it were to run on the global batch. For example, we want the step counter to increase only after all mini-batches passed, so the learning rate will be modified at the correct rate.
We want all the updates to take place only after all the mini-batches have passed. Technically, this means we want the updates to occur at the last step of the accumulation — when last is True.
So, if we could have just called the original optimizer’s get_updates(), while (1) making it use the accumulated gradients, and (2) causing all the variable updates to take place only at the last steps of the accumulation, we would have achieved what we wanted.
Fortunately, replacing (hooking) methods is really easy in Python, and by replacing a few methods with a different implementation we can easily achieve exactly that.
Optimizers call their get_gradients() from get_updates() to calculate the gradients of the parameters with respect to the loss. Therefore, we replace the optimizer’s get_gradients() with a function that does nothing but returning the accumulated gradients (agrads — the tensors we generated in line 13). This will cause the original optimizer to refer to the accumulated gradients in its algorithm and will solve (1). Let’s take a look at a simplified implementation for such a replacement method:
Regarding (2), variables in Keras can be assigned using 3 methods: K.update(), K.update_add(), and K.update_sub(). Optimizers use these methods for all updates — for the model parameters as well as for the optimizer parameters. We replace all three of them (can be seen on GitHub). We want all tensors created using those methods to assign values only in the last mini-batch and to do nothing otherwise. Therefore, in our methods — that replace the three — we wrap every value being assigned with a conditional switch and pass this switch to the respective method. If this is the last mini-batch (last is True), we assign the actual value to the variable, otherwise, we assign a value that does not affect the variable. For K.update_add() and K.update_sub() we assign zero, causing the variable to not actually increase or decrease. For K.update() we assign the current value of the variable, causing the variable to keep its current value. Let’s take a look at a simplified implementation for such replacement methods:
Back to our get_updates(), in lines 15–18 we actually replace all those methods. We use helper classes — subclasses of runai.utils.Hook — to do so.
In line 19 we call the original optimizer’s get_updates(). With all those methods replaced, we (1) make it refer to the accumulated gradients, and (2) cause all updates (“Assign” ops) to take place only when last is True.
We have two more things we have to do at the end of every step.
First, we have to update our variables to hold the current value of the accumulated gradients. This is done in line 33.
Second, we have to advance our step counter — self.iterations (line 36). To make sure this happens at the end of the step we generate the tensors under control dependencies of all other “Assign” ops. This causes the step counter update to take place only after all other updates have already taken place.
This article along with previous ones was meant to describe a few things in details:
The problems in batch sizing one might encounter while training neural networks and being limited by GPU memory.Why and how gradient accumulation might help in fixing them.A deep-dive into Keras optimizers and the things that need to be taken into consideration while building a gradient accumulation mechanism.How to use the Run:AI Python library and easily add gradient accumulation support to your Keras models.
The problems in batch sizing one might encounter while training neural networks and being limited by GPU memory.
Why and how gradient accumulation might help in fixing them.
A deep-dive into Keras optimizers and the things that need to be taken into consideration while building a gradient accumulation mechanism.
How to use the Run:AI Python library and easily add gradient accumulation support to your Keras models.
We hope we were able to shed some light on those subjects and help you understand them a bit better.
An open-source gradient accumulation tool, as well as usage examples and more resources, are available on GitHub.
Please share with us if you found this interesting and helpful in running your own neural networks!