Embedding Preprocessing in Keras Models with Normalization, StringLookup, and CategoryEncoding

Problem: preprocessing steps drift between training and serving

When you build a model that expects cleaned, scaled, or encoded inputs, it is common to write separate scripts that transform raw data before feeding it to the network. If the transformation logic changes or is omitted during deployment, the model sees data with a different distribution, leading to silent performance drops. This mismatch is known as training‑serving skew.

Thesis: Keras preprocessing layers let you bake transformations into the model graph

Keras provides layers such as Normalization, StringLookup, and CategoryEncoding that operate like any other layer. Because they are part of the model, they are saved with model.save() and executed identically during training, evaluation, and inference. Using them removes the need for external preprocessing pipelines and guarantees that the same statistics and mappings are applied everywhere.

Worked example: a model that accepts raw numeric and categorical columns

Assume a CSV file with two numeric features (age, salary) and one categorical feature (department). The goal is to predict a binary label.

1. Prepare a tf.data dataset that yields raw tensors

import tensorflow as tf
import pandas as pd

# Load data (replace with your own path)
df = pd.read_csv('your_data.csv')
# Features and label
features = df[['age', 'salary', 'department']]
label = df['target']

# Create a dataset of raw tensors
raw_ds = tf.data.Dataset.from_tensor_slices((dict(features), label))
raw_ds = raw_ds.shuffle(1000).batch(32)

Run the snippet in a terminal or notebook where you have installed TensorFlow >= 2.4 (e.g., pip install tensorflow>=2.4). No special permissions are required beyond those needed to install packages.

2. Instantiate and adapt preprocessing layers

# Numeric normalization
normalizer = tf.keras.layers.Normalization()
normalizer.adapt(raw_ds.map(lambda x, y: tf.stack([x['age'], x['salary']], axis=-1)))

# String lookup for department
lookup = tf.keras.layers.StringLookup(vocabulary=df['department'].unique())
# Category encoding – one‑hot
encoder = tf.keras.layers.CategoryEncoding(num_tokens=lookup.vocab_size(), output_mode='one_hot')

The adapt() call computes the mean and variance for numeric features and builds the vocabulary for the string lookup. Forgetting this step leaves the layers uninitialized, which causes the model to treat all inputs as zeros or unknown tokens.

3. Build the model with preprocessing layers as the first block

# Input layers for raw fields
age_in = tf.keras.Input(shape=(1,), name='age', dtype='tf.float32')
salary_in = tf.keras.Input(shape=(1,), name='salary', dtype='tf.float32')
dept_in = tf.keras.Input(shape=(1,), name='department', dtype='tf.string')

# Normalize numeric features
numeric = tf.keras.layers.Concatenate()([age_in, salary_in])
numeric_norm = normalizer(numeric)

# Encode categorical feature
encoded_dept = encoder(lookup(dept_in))

# Combine processed features
processed = tf.keras.layers.Concatenate()([numeric_norm, encoded_dept])

# Example dense head
x = tf.keras.layers.Dense(16, activation='relu')(processed)
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)

model = tf.keras.Model(inputs=[age_in, salary_in, dept_in], outputs=out)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

The model summary will show the preprocessing layers as part of the graph.

4. Train, save, and verify consistency

# Train
model.fit(raw_ds, epochs=5)

# Save the whole model (preprocessing included)
model.save('raw_input_model')

# Reload
reloaded = tf.keras.models.load_model('raw_input_model')

# Predict on a single raw sample (no external preprocessing)
sample = {
    'age': tf.constant([[30.0]]),
    'salary': tf.constant([[50000.0]]),
    'department': tf.constant([['Engineering']])
}
pred_original = model.predict(sample)
pred_reloaded = reloaded.predict(sample)

# The two predictions should be identical up to floating‑point rounding
print('Original:', pred_original)
print('Reloaded:', pred_reloaded)

You can verify that the preprocessing is embedded by inspecting the reloaded model’s layers (reloaded.layers) and confirming that Normalization, StringLookup, and CategoryEncoding appear before the dense layers.

Trade‑off and limitation

Embedding preprocessing simplifies deployment but couples the model to a specific data schema. If the raw input format changes (e.g., a new categorical value appears), you must retrain or at least re‑adapt the lookup layers, because the vocabulary is fixed after adapt(). Additionally, stateful layers like Normalization increase the model size slightly because they store statistics (mean, variance). For very large vocabularies, the StringLookup table can become a memory consideration; in such cases you may still opt for external hashing or embedding lookups.

Actionable closing

Start by identifying the numeric and categorical fields in your dataset. Create a Normalization for numeric columns, adapt it on a representative batch, and pair StringLookup with CategoryEncoding for each categorical column. Stack these layers as the first block of your Keras model, train as usual, and save the model with model.save(). Verify consistency by comparing predictions from the original and reloaded models on raw inputs. This approach eliminates a common source of bugs and makes your serving code as simple as calling model.predict(raw_batch).