How to build a wide-and-deep model using Keras in TensorFlow 2.0

Using Feature columns with the Keras Functional API

In TensorFlow 2.0, Keras has support for feature columns, opening up the ability to represent structured data using standard feature engineering techniques like embedding, bucketizing, and feature crosses. In this article, I will first show you a simple example of using the Functional API to build a model that uses features columns. Then, I will update the code to build out a full wide-and-deep model.

To keep the article short, I am focusing on just the model, but you can see the full notebook (including reading data using tf.data) on GitHub.

1. Install TensorFlow 2.0 if necessary

The code here assumes you are using TensorFlow 2.0. pip install it in Colab using:

!pip install -q tensorflow==2.0.0-alpha0`

or use Cloud AI Platform (CAIP) Notebooks to get a TensorFlow 2.0 Jupyter instance:

Use TensorFlow 2.0 to try out this code

2. Create Feature Columns

Create feature columns for each of the fields in your input:

real = {
    colname : fc.numeric_column(colname) \
          for colname in \
            (’dep_delay,taxiout,distance’).split(’,’)
}
sparse = {
      'origin' : fc.categorical_column_with_hash_bucket(’origin’, hash_bucket_size=1000),
      'dest'   : fc.categorical_column_with_hash_bucket(’dest’, hash_bucket_size=1000)
}

Here, I’m creating numeric columns for the float fields and hashed-categorical columns for the categorical fields.

3. Create Input to the model

For each of the inputs, also create a Keras Input layer, making sure to set the dtype and name for each of the input fields:

inputs = {
    colname : tf.keras.layers.Input(name=colname, shape=(), dtype='float32') \
          for colname in real.keys()
}
inputs.update({
    colname : tf.keras.layers.Input(name=colname, shape=(), dtype='string') \
          for colname in sparse.keys()
})

4. Do embedding, one-hot-encoding, etc.

In order to use a categorical variable in a deep learning model, we have to encode it. A sparse variable will have to be either embedded or one-hot encoded. So, let’s do both:

embed = {
       'embed_{}'.format(colname) : fc.embedding_column(col, 10) \
          for colname, col in sparse.items()
}
real.update(embed)
# one-hot encode the sparse columns
sparse = {
    colname : fc.indicator_column(col) \
          for colname, col in sparse.items()
}

5. Build a DNN model using Functional API

The key thing is to create a DenseFeatures layer to transform the inputs using the feature columns:

deep = tf.keras.layers.DenseFeatures(real.values())(inputs)
deep = tf.keras.layers.Dense(64, activation='relu')(deep)
deep = tf.keras.layers.Dense(16, activation='relu')(deep)
output = tf.keras.layers.Dense(1, activation='sigmoid')(deep)
model = tf.keras.Model(inputs, output)

That’s it! Call model.fit() etc. as usual.

6. Build a wide-and-deep network using Keras Functional API

If you want to build a wide-and-deep network, you want to wire the sparse features directly to the output node, but pass the real features through a set of dense layers. Here’s a model architecture that will do that:

def wide_and_deep_classifier(inputs, linear_feature_columns, dnn_feature_columns, dnn_hidden_units):
    deep = tf.keras.layers.DenseFeatures(dnn_feature_columns)(inputs)
    for numnodes in dnn_hidden_units:
        deep = tf.keras.layers.Dense(numnodes, activation='relu')(deep)        
    wide = tf.keras.layers.DenseFeatures(linear_feature_columns)(inputs)
    both = tf.keras.layers.concatenate([deep, wide])
    output = tf.keras.layers.Dense(1, activation='sigmoid')(both)
    model = tf.keras.Model(inputs, output)
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model
model = wide_and_deep_classifier(inputs, sparse.values(), real.values(), [64, 16])

Of course, this being Keras, you can easily sprinkle in dropout, batch normalization, etc. into the model.

Enjoy!

Resources:

1. See the full notebook on GitHub.
2. Read my book on doing Data Science on GCP