22 Nov 2017
Tensorflow Guide: Exponential Moving Average for Improved Classification
Parameter Selection via Exponential Moving Average
When training a classifier via gradient decent, we update the current classifier’s parameters $\theta$ via
where $\theta_t$ is the current state of the parameters and $\Delta \theta_t$ is the update step proposed by your favorite optimizer. Often times, after $N$ iterations, we simply stop the optimization procedure (where $N$ is chosen using some sort of decision rule) and use $\theta_N$ as our trained classifier’s parameters.
However, we often observe empirically that a post-processing step can be applied to improve the classifier’s performance. Once such example is Polyak averaging. A closely related—and quite popular—procedure is to take an exponential moving averaging (EMA) of the optimization trajectory $(\theta_n)$,
where $\lambda \in [0, 1)$ is the decay rate or momemtum of the EMA. It’s a simple modification to the optimization procedure that often yields better generalization than simply selecting $\theta_N$, and has also been used quite effectively in semi-supervised learning.
Implementation-wise, the best to apply EMA to a classifier is to use the built-in
tf.train.ExponentialMovingAverage function. However, the documentation doesn’t provide a guide for how to cleanly use
tf.train.ExponentialMovingAverage to construct an EMA-classifier. Since I’ve been playing with EMA recently, I thought that it would be helpful to write a gentle guide to implementing an EMA-classifier in Tensorflow.
For those who wish to dive straight into the full codebase, you can find it here. For self-containedness, let’s start with the code that constructs the classifier.
Here, I use a fairly standard CNN architecture. The first thing to note is the use of variable scoping. This puts all of the classifier’s variables within the scope
class/. To create the classifier, simply call
Once the classifier is created in the computational graph, variable scoping allows for easy access of the classifier’s trainable variables via
After getting the list of trainable variables via
tf.get_collection, we use
ema.apply, which serves two purposes. First, it constructs an auxiliary variable for each corresponding variable in
var_class to hold the exponential moving average. Next, it returns a tensorflow Op which updates the EMA variables. The object
ema can then access the EMA via the function
Populating the Classifier with the EMA Variables
So far, we’ve figured out how to create the EMA variables and how to access them. But what’s the easiest way to make the classifier use the EMA variables? Here, we leverage the
custom_getter argument that appears in
tf.variable_scope. According to documentation, whenever you call
tf.get_variable, the default getter gets an existing tensor according to the tensor variable’s name. However, a custom getter can be applied to change the existing tensor that is returned by
To construct the custom getter, locally define
ema_getter after you’ve already created the
To apply the EMA classifier during test-time, we simply call
classifier again, this time with the custom
And that’s it! We can now verify that applying EMA does in fact improve the performance of the classifier on the CIFAR-10 test data set.
You can find the full code for training the CIFAR-10 classifier below.
End of post