Up until now Iā€™ve been manually adjusting hyperparameters on my ML models until Iā€™m happy with the result. But now that Iā€™ve discovered how to automate it, Iā€™m never going back šŸŽ‰

In this post Iā€™ll show you how to run automated hyperparameter tuning on AI Platform via XGBoost model code packaged in a custom container. If you havenā€™t hit buzzword capacity yet, read on.

You might be wondering: what is a hyperparameter? And why do you need to tune it? Good question! Before I explain how to automate the tuning, let me quickly explain what hyperparameters are.

Hyperparameter tuning: the basics

When youā€™re building machine learning models, there are all sorts of configuration options you need to set before you start training. Things like the number of data points your model will process at once, the size of each layer, how much your model should adjust itā€™s learning rate, and the number of times it should iterate over all of your data. There are some rough guidelines on how to set these numbers (start with a few, small layers and build from there) but sometimes the process of setting them can seem arbitrary. These values are called hyperparameters.

Confession time: my typical workflow thus far has involved choosing hyperparameters that worked for someone elseā€™s model on a similar task and then haphazardly adjusting those values up and down depending on the results until I was happy with the accuracy (and of course, the results of my interpretability analysis). This approach works, but itā€™s fairly tedious and there may be a combination of hyperparameters I never discover.

The tools

šŸ’« Automated hyperparameter tuning to the rescue. Here are the tools Iā€™ll be using to show you how this works:

  • Dataset: Iā€™ll train a model using a subset of the NOAA weather data in BigQuery public datasets.
  • ML framework: To keep the model code short and sweet Iā€™ll build it with XGBoost
  • Containers: Iā€™ll be containerizing the model code with Docker and hosting it on Google Container Registry
  • Training & tuning: Finally, for training and hyperparam tuning Iā€™m using Cloud AI Platform, specifically making use of custom container training and atuomated hyperparameter tuning

Will it rain in the UK today?

Iā€™ve been visiting the UK a lot recently (itā€™s one of my favorite places šŸ˜€) and I thought it would be fun see if I could predict whether it would rain on a given day based on historical weather data. Thereā€™s an awesome public weather dataset available in BigQuery with weather data from all over the world going back to 1932.

To grab all weather data from the UK (filtered by lat/lng) I ran the following SQL query:

SELECT
  stn, year, mo, temp, dewp, stp, visib,
  wdsp, sndp, fog, b.lat, b.lon,
  cast(b.elev as float64) as elev, gust,
  if(prcp > 0 and rain_drizzle = "1", 1, 0) as rained
FROM `bigquery-public-data.noaa_gsod.gsod*` a
JOIN(
  SELECT usaf, lat, lon, elev FROM `bigquery-public-data.noaa_gsod.stations`
) b
ON a.stn = b.usaf
WHERE lat > 50.1 AND lat < 60.1
AND lon > -7.6 AND lon < 1.7

I did some additional filtering to remove null values and the result is 3,881 rows of data. Hereā€™s a preview:

UK weather data

With that weā€™re ready to build the model. And because this post doesnā€™t have any pictures, hereā€™s one of me in the UK on a day that couldnā€™t decide if it was sunny or rainy:

Sara with an umbrella

Building the XGBoost model

Before we containerize things, letā€™s setup the data and write the XGBoost model code. First weā€™ll read the data in as a Pandas DataFrame and split it into train and test sets:

import os
import pandas as pd
import xgboost as xgb

from sklearn.utils import shuffle

data = pd.read_csv(os.getcwd() + '/trainer/rain_uk.csv')
data = shuffle(data, random_state=2)

labels = data['rained']
data = data.drop(columns=['rained', 'year', 'stn'])

train_size = int(len(data) * .8)
train_data = data[:train_size]
test_data = data[train_size:]

train_labels = labels[:train_size]
test_labels = labels[train_size:]

I used os.getcwd() to get the filepath so that our script will be able to access it from within the container (which weā€™ll create soon). Our model will be a binary classifier, predicting 0 for no rain and 1 for rain.

The code to create our XGBClassifier and train it is simple. There are a lot of optional parameters we could pass in, but for now weā€™ll use the defaults (weā€™ll use hyperparameter tuning magic later to find the best values):

bst = xgb.XGBClassifier(
    objective='binary:logistic'
)

bst.fit(
    train_data.values, 
    train_labels, 
    eval_metric='error', 
    eval_set=[(test_data.values, test_labels)]
)

Then we can print the modelā€™s accuracy as a percentage, along with the error value for the last training epoch:

from sklearn.metrics import accuracy_score

y_pred = bst.predict(test_data.values)
predictions = [round(value) for value in y_pred]
accuracy = accuracy_score(test_labels, predictions)
print("Accuracy: %.2f%%" % (accuracy * 100.0))

error = list(bst.evals_result_.values())[0]['error'][-1]
print('Error: ', error)

On the first attempt at training this I got an accuracy of 86% (your exact % may vary) - not bad. The last thing weā€™ll want to do is save our model to Google Cloud Storage. First, use argparse to pass our Cloud Storage bucket as a command line argument:

import argparse

def get_args():
    parser = argparse.ArgumentParser(description='XGB on custom containers')
    parser.add_argument(
        '--model-dir',
        default=None,
        help='The directory to store the model')
    args = parser.parse_args()
    return args

And then invoke subprocess to save the model file via gsutil:

import subprocess

# Save the model locally
bst.save_model('model.bst')

# Save to GCS
if args.model_dir:
    tmp_model_file = os.getcwd() + '/model.bst'
    subprocess.check_call([
        'gsutil', 'cp', tmp_model_file,
        os.path.join(args.model_dir, 'model.bst')])

Using Cloud hyperparameter tuning

Time to add the hyperparameter tuning code! For that weā€™ll be using the hypertune package made available by Google Cloud. To use it there are a few steps:

  • Choose the hyperparameters you want to tune
  • Create a config.yaml file with the guidelines for tuning those parameters
  • Add some hypertune code to your model script to use those hyperparameters in your training
  • Submit the training job to Cloud AI Platform

Which parameters will we tune?

You can see the full list of parameters for XGBClassifier in the docs. Weā€™ll experiment with two:

  • max_depth: the maximum tree depth for our model, default is 3 (remember XGB is a boosted tree framework, more on that here)
  • learning_rate: how much our modelā€™s weights will change every time theyā€™re updated, default is 0.1

Now that weā€™ve chosen the hyperparameters we want Cloud to find the optimal values for, we can create a config file with some guidelines for tuning. Remember above when we calculated the error of our trained XGB model? Thatā€™s a number we want to minimize. Weā€™ll pass this to the hypertune service, so it has a metric for judging the success of each training run (called a trial).

Hereā€™s what our config file looks like:

trainingInput:
  hyperparameters:
    goal: MINIMIZE
    hyperparameterMetricTag: "error"
    maxTrials: 10
    maxParallelTrials: 2
    enableTrialEarlyStopping: True
    params:
    - parameterName: max-depth
      type: INTEGER
      minValue: 3
      maxValue: 10
    - parameterName: learning-rate
      type: DOUBLE
      minValue: 0.01
      maxValue: 0.9

Itā€™s pretty easy to understand, but hereā€™s a bit more detail on the maxParallelTrials config: hypertune can run multiple trials in parallel. The benefit of this is that your job will finish faster but the service can only optimize values based on data from completed trials, so this is a tradeoff youā€™ll want to consider.

Adding hypertune code

To add our hyperparameters to our model code, weā€™ll parse them as command line arguments by adding them to the get_args() function defined above:

parser.add_argument(
    '--max-depth',
    help='Max depth of XGB tree',
    default=3
)
parser.add_argument(
    '--learning-rate',
    help='Learning rate of XGB model',
    default=0.1
)

Note that the names of the parameters here correspond with what weā€™ve defined in our config file. Weā€™ll call this function at the beginning of our main method:

def main():
  args = get_args()
  # Data + model code goes here..

Next we want to send the hypertune metrics to Cloud. Weā€™ll add some code to do that after training and calculating evaluation metrics:

hpt = hypertune.HyperTune()
hpt.report_hyperparameter_tuning_metric(
    hyperparameter_metric_tag='error',
    metric_value=error,
    global_step=1)

Since you donā€™t need to set the number of epochs for XGBoost Iā€™ve set the global_step to 1. The hyperparameter_metric_tag correpsonds to our config file.

Finally, add the args created above to your model training code:

bst = xgb.XGBClassifier(
    objective='binary:logistic',
    max_depth=args.max_depth,
    learning_rate=args.learning_rate
)

When we kick off our hyperparameter training job, itā€™ll change max_depth and learning_rate for each training trial.

Time to containerize šŸ“¦

I decided to containerize the model code to try out the custom container training feature on Cloud AI Platform (also because Iā€™ve never played with Docker and this seemed like a good excuse). Learning all these new tools was surprisingly frictionless.

You donā€™t have to use a custom container to train an XGBoost model on AI Platform, but Iā€™ve used it here to show you how it works. And the cool thing about custom containers is that you can write an ML pipeline using whatever framework youā€™d like and itā€™ll train on AI Platform. The possibilities are endless.

So far weā€™ve created a single Python file with our model code and config.yaml with our hypertune configs. Weā€™ll now package up the code:

  • Dockerfile (weā€™ll write this next)
  • config.yaml
  • trainer/
    • model.py
    • rain_uk.csv (our data file)

Writing the Dockerfile

In our Dockerfile weā€™ll include all the commands needed to run our image. Itā€™ll install all the libraries weā€™re using and set up the entry point for our training code.

You can find the Dockerfile I used for this here.

Pushing to Google Container Registry

To use our container we need to first build it. If you havenā€™t already, youā€™ll want to run gloud auth configure-docker to connect gcloud with your Docker credentials.

Then youā€™ll build your container passing it the url of your image in Google Container Registry (details on the format below):

export IMAGE_URI=gcr.io/your-gcp-project/your-repo-name:image_tag
docker build -f Dockerfile -t $IMAGE_URI

You can choose whatever youā€™d like for the image repo name and tag. Next, test the container by running it locally:

docker run $IMAGE_URI

And push it to Container Registry:

docker push $IMAGE_URI

With that weā€™ve set up everything we need to start training (and automated tuning, obvs šŸ”§).

Let the automated hyperparameter party begin

You can kick off the training job with a single gcloud command. Update this with the name of your training job (can be anything you want), the region youā€™d like to run it in, your image URI, and the storage bucket to write model artifacts to:

gcloud ai-platform jobs submit training $JOB_NAME \
  --scale-tier BASIC \
  --region $REGION \
  --master-image-uri $IMAGE_URI \
  --config config.yaml \
  -- \
  --model-dir="gs://$BUCKET_NAME/$MODEL_DIR"

Next head over to the AI Platform > Jobs section of your console and select the job you just kicked off. In this job details view, you should see a table where your hypertune summary will be logged as trials complete:

HyperTune summary

You can refresh this page as training continues to see the results of the latest trial. In the screenshot above Iā€™m sorting by lowest error (the thing Iā€™m trying to minimize). You can also click View logs to monitor your job as it trains.

Looks like my best trial was #4, using a max_depth of 10 and a learning rate of 0.362. I updated my model to use those parameters, then re-ran training and accuracy increased to 90% - quite an improvement! There are obviously many more hyperparameters I could tune, which might improve my accuracy even more.

It would have taken ages for me to find this combination of hyperparams on my own, which is the best part of automated hyperparameter tuning. Thereā€™s also a lot of magic happening under the hood in how hypertune updates your hyperparameters after each trial. It does this using Bayesian optimization, which I wonā€™t pretend to be an expert in so instead Iā€™ll just link you to this blog post.

Learn more

If youā€™ve made it this far in the post, tweet me ā€œIā€™m hyped about hypertune!ā€ In all seriousness, here are some useful resources if you want to start using these tools on your own models:

Iā€™m always open to feedback and suggestions for future posts (and cool datasets, because there are never enough). Let me know what you think on Twitter at @SRobTweets.