Building Powerful CNNs with the TensorFlow Keras Functional API

Convolutional neural networks (CNNs) have revolutionized the field of computer vision, powering state-of-the-art models for complex visual understanding tasks like image classification, object detection, and semantic segmentation. CNNs are designed to process grid-like data such as images by learning a hierarchy of visual features through the use of convolution filters.

When building CNNs with the popular Keras deep learning library, many tutorials and code examples use the Sequential API, which allows you to create models layer-by-layer in a linear stack. While the Sequential API makes it simple to build standard feedforward networks, it has limitations in supporting more advanced CNN architectures that have complex branching and merging of layers.

This is where the Keras Functional API comes to the rescue. The Functional API provides a more flexible way to build sophisticated CNN models by allowing you to define the model as a directed acyclic graph (DAG) of layers. With the Functional API, you can easily create models that have:

• Multiple input branches that process different types of data (e.g. images and categorical features)
• Multiple output heads for multi-task learning
• Skip connections between non-adjacent layers (e.g. residual networks)
• Shared layers that are reused in different parts of the model

In this guide, we‘ll walk through a complete example of how to build a CNN for image classification using the Keras Functional API with the TensorFlow 2.0 backend. We‘ll be working with the classic MNIST handwritten digits dataset.

Let‘s dive in!

Step 1: Loading and Preprocessing the Data

The first step is to load the MNIST dataset and preprocess the images and labels into a format suitable for training a CNN. We‘ll use the mnist dataset utility in tf.keras.datasets to load the data as NumPy arrays split into training and test sets.

We also need to preprocess the data by:

1. Reshaping the images from (60000, 28, 28) to (60000, 28, 28, 1) for the CNN input layer
2. Normalizing the pixel values from [0, 255] to [0, 1]
3. One-hot encoding the labels

Here‘s the code:

from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

# Load the MNIST dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# Reshape images to 4D array with shape (num_samples, height, width, channels) 
x_train = x_train.reshape((60000, 28, 28, 1))
x_test = x_test.reshape((10000, 28, 28, 1))

# Normalize pixel values to [0, 1]
x_train = x_train.astype(‘float32‘) / 255
x_test = x_test.astype(‘float32‘) / 255

# One-hot encode the labels
num_classes = 10
y_train = to_categorical(y_train, num_classes)
y_test = to_categorical(y_test, num_classes)

Step 2: Defining the CNN Architecture

Next, we‘ll use the Keras Functional API to define the architecture of our CNN model. The key idea is to create an Input layer that specifies the shape of the input data, then connect it to a stack of convolutional (Conv2D), pooling (MaxPooling2D), and fully-connected (Dense) layers to map the input to the output predictions.

Here‘s a simple CNN architecture suitable for MNIST digit classification:

from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.models import Model

# Input layer
inputs = Input(shape=(28, 28, 1))

# Convolutional layers
x = Conv2D(32, kernel_size=(3, 3), activation=‘relu‘)(inputs)
x = MaxPooling2D(pool_size=(2, 2))(x)
x = Conv2D(64, kernel_size=(3, 3), activation=‘relu‘)(x)
x = MaxPooling2D(pool_size=(2, 2))(x)

# Flatten feature maps
x = Flatten()(x)

# Fully-connected layers for classification
x = Dense(128, activation=‘relu‘)(x)
x = Dropout(0.5)(x)
outputs = Dense(num_classes, activation=‘softmax‘)(x)

# Create model by specifying inputs and outputs
model = Model(inputs=inputs, outputs=outputs)

Let‘s break this down:

• We first create an Input layer that specifies the shape of input images as (28, 28, 1)
• We then apply two Conv2D layers with 32 and 64 filters respectively, using a 3×3 kernel and ReLU activation. Each Conv2D layer learns to detect visual features in the input.
• We apply MaxPooling2D after each convolutional layer to downsample the feature maps and add translation invariance
• We flatten the 2D feature maps into a 1D vector using Flatten to pass to the fully-connected layers
• We apply a Dense layer with 128 units and ReLU activation for learning high-level representations
• We apply Dropout for regularization to reduce overfitting
• Finally, we apply the output Dense layer with 10 units (one per class) and softmax activation to get the predicted probabilities for each class

We create the model by passing the inputs and outputs to Model(), which initializes a model instance that we can train and make predictions with.

Note how the Functional API makes it easy to connect layers in a graph structure by specifying the inputs argument for each layer. The output of one layer can be passed as the input to multiple layers.

We can visualize the CNN architecture by calling plot_model():

from tensorflow.keras.utils import plot_model

plot_model(model, to_file=‘cnn_model.png‘, show_shapes=True)

Step 3: Training the Model

With our model architecture defined, we‘re ready to train it on the MNIST dataset. First, we need to configure the model for training by specifying the loss function, optimizer, and metrics to monitor:

model.compile(loss=‘categorical_crossentropy‘,
              optimizer=‘adam‘, 
              metrics=[‘accuracy‘])

We use the categorical cross-entropy loss since we have a multi-class classification problem, and the Adam optimizer which is a popular choice that adapts the learning rate during training. We‘ll monitor the accuracy metric to evaluate the model.

To start the training process, we simply call fit() on our model, passing in the training data, number of epochs, and batch size:

batch_size = 128
epochs = 10

history = model.fit(x_train, y_train,
                    batch_size=batch_size,
                    epochs=epochs,
                    validation_split=0.1)

This will train the model for 10 epochs using a batch size of 128 samples. We set aside 10% of the training data as a validation set to monitor overfitting during training. The history object returned contains the loss and accuracy metrics tracked over the course of training.

Here‘s the typical output during training:

Epoch 1/10
422/422 [==============================] - 18s 43ms/step - loss: 0.2432 - accuracy: 0.9265 - val_loss: 0.0613 - val_accuracy: 0.9807
Epoch 2/10
422/422 [==============================] - 17s 40ms/step - loss: 0.0674 - accuracy: 0.9792 - val_loss: 0.0415 - val_accuracy: 0.9858
...
Epoch 10/10
422/422 [==============================] - 18s 42ms/step - loss: 0.0291 - accuracy: 0.9913 - val_loss: 0.0256 - val_accuracy: 0.9925

After 10 epochs, our simple CNN model reaches about 99% validation accuracy, indicating that it has learned to accurately classify the handwritten digits.

We can plot the training curves to visualize how the model‘s performance improved over time:

import matplotlib.pyplot as plt

acc = history.history[‘accuracy‘]
val_acc = history.history[‘val_accuracy‘]
loss = history.history[‘loss‘]
val_loss = history.history[‘val_loss‘]

epochs_range = range(1, epochs+1)

plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label=‘Training Accuracy‘)
plt.plot(epochs_range, val_acc, label=‘Validation Accuracy‘)
plt.legend(loc=‘lower right‘)
plt.title(‘Training and Validation Accuracy‘)

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label=‘Training Loss‘) 
plt.plot(epochs_range, val_loss, label=‘Validation Loss‘)
plt.legend(loc=‘upper right‘)
plt.title(‘Training and Validation Loss‘)
plt.show()

The training and validation curves for accuracy increase while the loss curves decrease, indicating the model is learning effectively without overfitting.

Step 4: Evaluating the Model

Once our model is trained, we can evaluate its final performance on the test set, which the model did not see during training:

test_loss, test_acc = model.evaluate(x_test, y_test)
print(‘Test accuracy:‘, test_acc)

Running this gives:

Test accuracy: 0.9915

Our model achieves over 99% accuracy in classifying the handwritten digits in the test set, confirming its strong predictive performance.

We can also use the trained model to make predictions on new images:

predictions = model.predict(x_test)

The predict() method returns the predicted class probabilities for each image in x_test. We can get the predicted class labels by taking the argmax() of the probabilities:

predicted_labels = np.argmax(predictions, axis=1)

Improving the Model

There are many strategies we can employ to improve our CNN model‘s performance and generalization ability:

• Data augmentation: We can artificially expand our training dataset by applying random transformations (rotations, translations, scaling, etc.) to the input images. This helps the model learn to be invariant to these transformations and reduces overfitting. Keras has an ImageDataGenerator class that automates this process.

• Regularization: In addition to dropout, we can apply other regularization techniques to constrain the model complexity and reduce overfitting, such as L1/L2 weight regularization and batch normalization. These are easily added as additional layers or arguments to existing layers in the Functional API.

• Hyperparameter optimization: The performance of deep learning models heavily depends on the choice of hyperparameters such as the learning rate, batch size, number and size of layers, etc. We can use techniques like grid search, random search, or Bayesian optimization to automatically find the optimal hyperparameter values that maximize model performance. Keras has a RandomSearch and Hyperband tuner classes for this purpose.

• Experiment with more advanced architectures: There are many state-of-the-art CNN architectures that have proven effective for image classification tasks, such as VGGNet, GoogLeNet, ResNet, and EfficientNet. We can implement these architectures using the Functional API and explore ensembling multiple models to boost performance.

Productionizing the Model

Once we have a trained CNN model that meets our performance requirements, we may want to deploy it as part of a production application to serve real-time predictions. Here are a few tips to make your model production-ready:

• Save the model: The trained Keras model can be saved to disk in the TensorFlow SavedModel format using model.save(‘path/to/model‘). This saves the model architecture, weights, and even the optimizer state so you can easily load the model later for inference or further training.

• Optimize for inference: To reduce latency and resource usage when serving your model, you can optimize it for inference by removing unnecessary nodes, quantizing weights, and fusing common operations using the TensorFlow Lite Converter (tf.lite.TFLiteConverter). This is especially important for deploying models on mobile/IoT edge devices.

• Create a prediction service: A common way to expose your trained model to other applications is to wrap it as a web service API endpoint. You can use a web framework like Flask or FastAPI to create a REST API that takes an input image, preprocesses it, runs the model inference, and returns the predicted class probabilities. Alternatively, you can use a serverless platform like Google Cloud Functions or AWS Lambda to deploy your model as a scalable and cost-effective prediction service.

• Monitor and maintain: Like any production system, it‘s important to continuously monitor your deployed model‘s performance and resource usage to catch any issues or drift in the input data distribution. You should also have a pipeline in place to retrain and update your model on new data to maintain its accuracy over time.

Conclusion

In this guide, we‘ve seen how to use the powerful Keras Functional API to build and train CNN models for image classification tasks. The Functional API provides the flexibility to define complex model architectures with multiple inputs/outputs, shared layers, and non-sequential topologies, which is essential for many real-world applications.

We walked through a complete example of building a CNN to classify handwritten digits from the MNIST dataset, including data preprocessing, model architecture definition, training, and evaluation. We also discussed various strategies to improve model performance and tips for deploying trained models in production.

To learn more, check out these helpful resources:

• The official Keras documentation on the Functional API: https://keras.io/guides/functional_api/
• A detailed guide to the Keras Functional API: https://www.tensorflow.org/guide/keras/functional
• Overview of CNN architectures for visual recognition: https://arxiv.org/abs/1605.07678
• Tutorials on productionizing Keras models: https://www.tensorflow.org/tfx/tutorials

I hope this guide has given you a solid foundation for building CNN models with the Keras Functional API. Happy deep learning!