Finding the Optimal Learning Rate for Your Neural Network with LearningRateScheduler in TensorFlow

ByJordan Brown

When training deep neural networks, the learning rate is one of the most important hyperparameters to tune. The learning rate controls the size of the updates made to the model‘s weights during optimization. Set the learning rate too high and the model may fail to converge or even diverge. Set it too low and training will progress very slowly.

Achieving the fastest possible convergence requires using the largest learning rate that still allows the model to converge reliably. However, the optimal learning rate can be difficult to determine ahead of time and may change over the course of training. Fortunately, TensorFlow provides a useful callback called LearningRateScheduler that lets you automatically adjust the learning rate during training according to a predefined schedule. Even better, you can use LearningRateScheduler to test out a range of learning rates in a single training run, allowing you to easily identify the optimal learning rate for your model and dataset.

In this post, we‘ll take an in-depth look at how to leverage LearningRateScheduler to find the best learning rate for training a neural network on time series data. We‘ll go through a complete example of generating a synthetic time series, preparing a windowed dataset, defining a model, and using LearningRateScheduler to select an optimal learning rate. Finally, we‘ll evaluate the trained model‘s performance on a held-out test set.

Let‘s get started!

Generating a Synthetic Time Series

Before we can train a model, we need some data. For this example, we‘ll use a simple synthetic time series generated from a combination of trend, seasonality, and noise components. Here‘s the code to generate the series:

def trend(time, slope=0):
    return time * slope

def seasonal_pattern(season_time):
    return np.where(season_time < 0.4,
                    np.cos(season_time * 2 * np.pi),
                    1 / np.exp(3 * season_time))

def seasonality(time, period, amplitude=1, phase=0):
    season_time = ((time + phase) % period) / period
    return amplitude * seasonal_pattern(season_time)

def noise(time, noise_level=1, seed=None):
    rnd = np.random.RandomState(seed)
    return rnd.randn(len(time)) * noise_level

time = np.arange(4 * 365 + 1, dtype="float32")
baseline = 10
amplitude = 20
slope = 0.09
noise_level = 5

series = baseline + trend(time, slope) + seasonality(time, period=365, amplitude=amplitude)
series += noise(time, noise_level, seed=42)

This generates a 4-year daily time series with an upward trend, yearly seasonal pattern, and some random noise. We can visualize the resulting series:

def plot_series(time, series, format="-", start=0, end=None):
    plt.plot(time[start:end], series[start:end], format)
    plt.xlabel("Time")
    plt.ylabel("Value")
    plt.grid(True)

plt.figure(figsize=(10, 6))
plot_series(time, series)
plt.show()

Preparing the Dataset

With our time series in hand, the next step is to prepare a dataset for training and evaluating the model. We‘ll split the data into training and validation sets, then use TensorFlow‘s Dataset API to create a windowed dataset suitable for training a feedforward neural network.

Here‘s a function that will split a univariate time series into a windowed dataset:

def windowed_dataset(series, window_size, batch_size, shuffle_buffer):
    dataset = tf.data.Dataset.from_tensor_slices(series)
    dataset = dataset.window(window_size + 1, shift=1, drop_remainder=True)
    dataset = dataset.flat_map(lambda window: window.batch(window_size + 1))
    dataset = dataset.shuffle(shuffle_buffer).map(lambda window: (window[:-1], window[-1]))
    dataset = dataset.batch(batch_size).prefetch(1)
    return dataset

This function takes the time series, desired window size, batch size, and a shuffle buffer size. It first creates a dataset containing individual time steps. It then applies a sliding window to extract input/output pairs of fixed window size from the series. The resulting dataset contains tuples of (input_window, output_value). Finally, it shuffles the windows and batches them up for training.

To prepare the training set, we‘ll take the first 1000 days of data and split it into a windowed dataset:

split_time = 1000
time_train = time[:split_time]
x_train = series[:split_time]

window_size = 20
batch_size = 32
shuffle_buffer = 1000

dataset = windowed_dataset(x_train, window_size, batch_size, shuffle_buffer)

Defining the Model

We can now define a simple feedforward neural network to train on this dataset. The model will take a window of the 20 previous time steps as input and predict the next value in the sequence.

model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation=‘relu‘, input_shape=[window_size]),
    tf.keras.layers.Dense(10, activation=‘relu‘),
    tf.keras.layers.Dense(1)
])

This model has two hidden layers with 10 units each and ReLU activation, plus a linear output layer.

Finding a Good Learning Rate with LearningRateScheduler

Now for the key part – using the LearningRateScheduler to find an optimal learning rate. The idea is to train the model for a few epochs while rapidly increasing the learning rate from a very small value to a very large one. By looking at the resulting loss curve, we can identify the point just before the loss starts to increase, which indicates that the learning rate has gotten too large for stable training.

Here‘s how we can implement this:

lr_schedule = tf.keras.callbacks.LearningRateScheduler(
    lambda epoch: 1e-8 * 10**(epoch / 20))

model.compile(loss="mse", optimizer=tf.keras.optimizers.SGD(1e-8, momentum=0.9))
history = model.fit(dataset, epochs=100, callbacks=[lr_schedule], verbose=0)

The LearningRateScheduler takes a function that maps an epoch number to a learning rate. In this case, we use an exponential schedule that starts at 1e-8 and increases by a factor of 10 every 20 epochs.

We compile the model with a starting learning rate of 1e-8 and train it for 100 epochs, passing the LearningRateScheduler as a callback. By setting verbose=0, we suppress the training output, since we‘re just using this run to find a good learning rate, rather than actually train the model.

After training, we can plot the loss versus learning rate to see where things start to break down:

lrs = 1e-8 * 10 ** (np.arange(100) / 20)
plt.figure(dpi=150)
plt.semilogx(lrs, history.history["loss"])
plt.axis([1e-8, 1e-3, 0, 300])
plt.xlabel("Learning Rate")
plt.ylabel("MSE Loss")
plt.show()

The loss decreases steeply until a learning rate of around 3e-6, then starts to increase again. This indicates that 3e-6 is approximately the largest usable learning rate for this dataset and model.

Training the Final Model

Now that we‘ve identified a good learning rate, we can train the model for real. We‘ll instantiate a new model with the same architecture, compile it with the selected learning rate, and train for 100 epochs.

model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation="relu", input_shape=[window_size]),
    tf.keras.layers.Dense(10, activation="relu"),
    tf.keras.layers.Dense(1)
])

model.compile(loss="mse", optimizer=tf.keras.optimizers.SGD(3e-6, momentum=0.9))
history = model.fit(dataset, epochs=100, verbose=1)

Evaluating Performance

Finally, let‘s see how well the trained model performs on the held-out portion of the series. The following code generates predictions for the validation period and plots them against the true values:

time_valid = time[split_time:]
x_valid = series[split_time:]

model.reset_states()
forecast = []
for time in range(len(x_valid)):
    forecast.append(model.predict(x_valid[time:time + window_size][np.newaxis]))
forecast = forecast[window_size:]

results = np.array(forecast)[:, 0, 0]

plt.figure(figsize=(10, 6))
plot_series(time_valid, x_valid)
plot_series(time_valid[window_size:], results)
plt.title("Validation Period Forecast")
plt.show()

As a final quantitative check, we can also compute the mean absolute error between the forecasts and the actual values: 

tf.keras.metrics.mean_absolute_error(x_valid[window_size:], results).numpy()

# MAE: 3.87

The model achieves a respectable MAE of 3.87 on the validation period, indicating that it has learned to forecast this simple synthetic series reasonably well.

Conclusions and Further Reading

In this post, we‘ve seen how to use the LearningRateScheduler callback in TensorFlow to find an optimal learning rate for training a neural network. By rapidly increasing the learning rate over a few epochs and plotting the resulting loss curve, we can quickly get a sense of the usable range of learning rates for a given model and dataset.

LearningRateScheduler is a simple but powerful tool to have in your deep learning toolkit. While there are more advanced techniques like cyclical learning rates and stochastic gradient descent with warm restarts, LearningRateScheduler remains an easy and effective way to select a good learning rate with minimal fuss.

To learn more, I recommend the following resources:

I hope you‘ve found this post helpful! Remember, choosing a good learning rate is key to fast, stable training of neural networks. LearningRateScheduler makes this process much easier. Give it a try on your own projects!

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts