Feature Engineering using custom Keras Layers for a complete training and inference pipeline.

Fernando Nieuwveldt
3 min readMar 15, 2022

In a previous post we looked at applying feature engineering with Keras Lambda layers. We also mentioned that Lambda layers have some serialization limitations. In this post we will see how we can overcome these issues by subclassing Keras Layer interface and create custom layers to achieve feature engineering.

Quick recap on the previous post:

We implemented a training and inference pipeline natively with Keras/Tensorflow. From loading data with tf.data. We utilized Lambda layers for feature engineering. We ended up with a training pipeline where feature engineering was part of the network architecture and can be persisted and loaded for inference as standalone. From a design and serving perspective this has some benefits if for example one used Tensorflow serving to serve models. Data does not need to be engineered with a different library before feeding it your serving endpoint.

Steps we will follow:

  • Load data with tf.data
  • Create Input layer
  • Subclass Keras Layer class
  • Train model

We will again use the heart disease dataset. Lets import tensorflow and read in the data:

Create a dictionary of Input objects for each feature:

We will look at two ways to create our Feature layer. Scenario 1 is using multiple custom Keras layers for the different features and Scenario 2 uses only one layer.

Scenario 1: Multiple custom Layers

Here we will create separate custom layers for the different features and as a final step we will concatenate our layers to create one complete feature layer.

Let us build up our graph using our custom layers.

Next we will create a Keras Model to aid in visualising our network graph. This will illustrate how everything flows together and showcases layer dependencies.

This will produce the following graph:

Scenario 1 has a bit of steps, but one ends up with an informative graph showing how all the inputs feeds into the network and how the features are created. In the next section we will use a single custom layer to create the same features as above. We will however not be able to plot the same graph as below as the feature layer will be contained in one.

Scenario 2: Single custom layer

In this section we will implement a simpler feature layer.

This is all we need if we use only one layer. Lets compile and build our model:

In conclusion we showed two ways to implemented custom layers for feature engineering. Scenario 1 used multiple layers and involved more steps to implement, but also produced an informative graph displaying the flow of features. Scenario 2 was quicker and had less steps to implement.

--

--

Fernando Nieuwveldt

I am an ML Engineer | Data scientist with interests in Deep learning and building systems and software for ML. https://www.linkedin.com/in/fernandonieuwveldt