Site Logo

Building Capsule Networks with TensorFlow

Capsule Networks (CapsNets), introduced by Geoffrey Hinton and his team in 2017, represent a paradigm shift in deep learning, aiming to address limitations of traditional Convolutional Neural Networks (CNNs). Unlike CNNs, which rely on scalar outputs and pooling operations, CapsNets use capsules—groups of neurons that output vectors to represent entities and their properties (e.g., position, orientation, or scale). This structure preserves spatial hierarchies and improves generalization, especially for tasks like image recognition where viewpoint invariance is crucial. In this blog, we’ll explore how to build CapsNets using TensorFlow, covering their architecture, implementation, and practical applications. The guide is designed to be comprehensive, providing detailed explanations and code examples to help you implement CapsNets effectively.

Understanding Capsule Networks

Capsule Networks aim to overcome CNN shortcomings, such as loss of spatial relationships due to pooling and poor handling of pose variations. A capsule is a group of neurons whose output vector represents the presence and properties of an entity (e.g., an object or part of an object). The length of the vector indicates the probability of the entity’s presence, while its orientation encodes properties like rotation or scale.

CapsNets use dynamic routing to determine how lower-level capsules contribute to higher-level capsules, allowing the network to learn hierarchical relationships. This is a departure from CNNs, which use max-pooling to reduce spatial dimensions, often discarding precise positional information.

Key components of CapsNets include: 1. Primary Capsules: Extract low-level features (e.g., edges or textures) and group them into vectors. 2. Digit Capsules: Represent high-level entities (e.g., digits in MNIST) with vectors encoding class presence and pose. 3. Dynamic Routing: An iterative process that routes lower-level capsule outputs to higher-level capsules based on agreement. 4. Squash Activation: A non-linear function that normalizes capsule vectors to maintain probability-like lengths.

For a foundational understanding, refer to the original CapsNet paper (Dynamic Routing Between Capsules) and TensorFlow Neural Networks Intro.

Setting Up the TensorFlow Environment

To build CapsNets, we’ll use TensorFlow 2.x for its flexibility in custom layer development and eager execution. Ensure you have TensorFlow installed, and optionally use a GPU for faster training.

Install the required packages:

pip install tensorflow==2.15.0

For environment setup, see Installing TensorFlow and Setting Up Conda Environment. For GPU support, check GPU Memory Optimization.

We’ll implement CapsNets for the MNIST dataset, a standard benchmark for image classification, to demonstrate their effectiveness.

Defining the Capsule Network Architecture

A typical CapsNet for MNIST consists of: 1. Convolutional Layer: Extracts initial features from input images. 2. Primary Capsule Layer: Groups convolutional outputs into capsules. 3. Digit Capsule Layer: Represents digit classes with high-level capsules. 4. Reconstruction Network: A decoder that reconstructs the input image from digit capsule outputs, acting as a regularizer.

Let’s break down each component and implement them in TensorFlow.

Convolutional Layer

The initial convolutional layer extracts low-level features, similar to a CNN. For MNIST (28x28 grayscale images), we apply convolutions to produce feature maps.

import tensorflow as tf

class ConvLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(ConvLayer, self).__init__()
        self.conv = tf.keras.layers.Conv2D(
            filters=256, kernel_size=9, strides=1, activation='relu', padding='valid'
        )

    def call(self, inputs):
        return self.conv(inputs)

This layer produces 256 feature maps of size 20x20 (since 28 - 9 + 1 = 20).

Primary Capsule Layer

The primary capsule layer reshapes convolutional outputs into capsules. Each capsule is an 8-dimensional vector, and we create multiple capsules per spatial location. For MNIST, we aim for 32 capsule types, each with 6x6 spatial locations (due to a subsequent convolution with stride 2).

class PrimaryCaps(tf.keras.layers.Layer):
    def __init__(self, num_capsules=32, dim_capsule=8):
        super(PrimaryCaps, self).__init__()
        self.num_capsules = num_capsules
        self.dim_capsule = dim_capsule
        self.conv = tf.keras.layers.Conv2D(
            filters=num_capsules * dim_capsule, kernel_size=9, strides=2, padding='valid'
        )

    def call(self, inputs):
        outputs = self.conv(inputs)  # Shape: (batch, 6, 6, num_capsules * dim_capsule)
        outputs = tf.reshape(outputs, [-1, outputs.shape[1] * outputs.shape[2] * self.num_capsules, self.dim_capsule])
        return self.squash(outputs)

    def squash(self, inputs):
        squared_norm = tf.reduce_sum(tf.square(inputs), axis=-1, keepdims=True)
        scale = squared_norm / (1.0 + squared_norm)
        return scale * inputs / tf.sqrt(squared_norm + 1e-8)

The squash function normalizes capsule vectors, ensuring their lengths represent probabilities (between 0 and 1). For more on activation functions, see Activation Functions.

Digit Capsule Layer with Dynamic Routing

The digit capsule layer represents each MNIST digit class with a 16-dimensional capsule. Dynamic routing ensures that lower-level capsules (from PrimaryCaps) contribute to higher-level capsules based on their agreement.

class DigitCaps(tf.keras.layers.Layer):
    def __init__(self, num_capsules=10, dim_capsule=16, num_routing=3):
        super(DigitCaps, self).__init__()
        self.num_capsules = num_capsules
        self.dim_capsule = dim_capsule
        self.num_routing = num_routing

    def build(self, input_shape):
        self.num_input_caps = input_shape[1]
        self.dim_input_caps = input_shape[2]
        self.W = self.add_weight(
            name='W',
            shape=[1, self.num_input_caps, self.num_capsules, self.dim_capsule, self.dim_input_caps],
            initializer='glorot_uniform',
            trainable=True
        )

    def call(self, inputs):
        inputs = tf.expand_dims(inputs, 2)  # (batch, num_input_caps, 1, dim_input_caps)
        inputs = tf.tile(inputs, [1, 1, self.num_capsules, 1])  # (batch, num_input_caps, num_capsules, dim_input_caps)
        W_tiled = tf.tile(self.W, [tf.shape(inputs)[0], 1, 1, 1, 1])
        u_hat = tf.matmul(W_tiled, inputs[..., tf.newaxis])  # (batch, num_input_caps, num_capsules, dim_capsule, 1)
        u_hat = tf.squeeze(u_hat, axis=-1)  # (batch, num_input_caps, num_capsules, dim_capsule)

        b = tf.zeros([tf.shape(inputs)[0], self.num_input_caps, self.num_capsules, 1])
        for i in range(self.num_routing):
            c = tf.nn.softmax(b, axis=2)  # Coupling coefficients
            s = tf.reduce_sum(c * u_hat, axis=1, keepdims=True)  # Weighted sum
            v = self.squash(s)  # Output capsules
            if i < self.num_routing - 1:
                b += tf.reduce_sum(u_hat * v, axis=-1, keepdims=True)
        return tf.squeeze(v, axis=1)  # (batch, num_capsules, dim_capsule)

    def squash(self, inputs):
        squared_norm = tf.reduce_sum(tf.square(inputs), axis=-1, keepdims=True)
        scale = squared_norm / (1.0 + squared_norm)
        return scale * inputs / tf.sqrt(squared_norm + 1e-8)

Dynamic routing iteratively updates coupling coefficients (c) based on the agreement between lower-level predictions (u_hat) and higher-level capsule outputs (v). For more on custom layers, see Custom Layers.

Reconstruction Network

The reconstruction network regularizes the CapsNet by reconstructing the input image from the digit capsule outputs. It takes the active capsule (corresponding to the predicted digit) and passes it through dense layers.

class ReconstructionNet(tf.keras.layers.Layer):
    def __init__(self, output_shape=(28, 28, 1)):
        super(ReconstructionNet, self).__init__()
        self.dense1 = tf.keras.layers.Dense(512, activation='relu')
        self.dense2 = tf.keras.layers.Dense(1024, activation='relu')
        self.dense3 = tf.keras.layers.Dense(tf.reduce_prod(output_shape), activation='sigmoid')
        self.output_shape = output_shape

    def call(self, inputs, mask):
        masked_inputs = inputs * tf.expand_dims(mask, -1)  # Mask inactive capsules
        flattened = tf.reduce_mean(masked_inputs, axis=1)  # Average active capsule
        h = self.dense1(flattened)
        h = self.dense2(h)
        reconstruction = self.dense3(h)
        return tf.reshape(reconstruction, [-1] + list(self.output_shape))

Building the CapsNet Model

Combine the layers into a complete CapsNet model. The model outputs both the digit capsule vectors (for classification) and the reconstructed image (for regularization).

class CapsNet(tf.keras.Model):
    def __init__(self):
        super(CapsNet, self).__init__()
        self.conv_layer = ConvLayer()
        self.primary_caps = PrimaryCaps()
        self.digit_caps = DigitCaps()
        self.reconstruction_net = ReconstructionNet()
        self.decoder_loss_weight = 0.0005

    def call(self, inputs, training=False):
        x = self.conv_layer(inputs)
        x = self.primary_caps(x)
        digit_caps = self.digit_caps(x)
        reconstructions = self.reconstruction_net(digit_caps, tf.norm(digit_caps, axis=-1))
        return digit_caps, reconstructions

    def margin_loss(self, y_true, y_pred):
        L = y_true * tf.square(tf.maximum(0.0, 0.9 - tf.norm(y_pred, axis=-1))) + \
            0.5 * (1 - y_true) * tf.square(tf.maximum(0.0, tf.norm(y_pred, axis=-1) - 0.1))
        return tf.reduce_mean(tf.reduce_sum(L, axis=1))

    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            digit_caps, reconstructions = self(x, training=True)
            margin_loss = self.margin_loss(y, digit_caps)
            reconstruction_loss = tf.reduce_mean(tf.square(reconstructions - x))
            total_loss = margin_loss + self.decoder_loss_weight * reconstruction_loss
        grads = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        return {'loss': total_loss, 'margin_loss': margin_loss, 'reconstruction_loss': reconstruction_loss}

The margin loss encourages correct capsules to have longer vectors (close to 0.9) and incorrect ones to have shorter vectors (close to 0.1). For more on loss functions, see Loss Functions.

Training the CapsNet

Let’s train the CapsNet on the MNIST dataset. Load and preprocess the data, then train the model.

# Load MNIST dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train[..., tf.newaxis].astype('float32') / 255.0
x_test = x_test[..., tf.newaxis].astype('float32') / 255.0
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

# Initialize and compile model
model = CapsNet()
model.compile(optimizer=tf.keras.optimizers.Adam(0.001))

# Train model
model.fit(x_train, y_train, batch_size=128, epochs=20, validation_data=(x_test, y_test))

Monitor training progress with TensorBoard (TensorBoard Training). For early stopping, see Early Stopping.

Evaluating the Model

Evaluate the model’s performance on the test set by computing classification accuracy based on the capsule vector lengths.

digit_caps, _ = model(x_test)
predictions = tf.argmax(tf.norm(digit_caps, axis=-1), axis=1)
accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, tf.argmax(y_test, axis=1)), tf.float32))
print(f"Test Accuracy: {accuracy:.4f}")

For more evaluation techniques, see Evaluating Performance.

Practical Applications and Extensions

CapsNets excel in tasks requiring spatial hierarchy awareness, such as:

  • Image Classification: Robust to rotations and transformations ([Image Classification](/tensorflow/computer-vision/image-classification)).
  • Object Detection: Preserve part-whole relationships ([Object Detection](/tensorflow/computer-vision/object-detection)).
  • Medical Imaging: Analyze structured data like MRI scans ([MRI Segmentation](/tensorflow/computer-vision/mri-segmentation)).

To extend CapsNets:

  • Optimize Performance: Use mixed precision training ([Mixed Precision](/tensorflow/fundamentals/mixed-precision)).
  • Deploy Models: Convert to TensorFlow Lite for mobile deployment ([TensorFlow Lite](/tensorflow/production/tensorflow-lite-mobile)).
  • Explore Variants: Experiment with attention-based capsules or hybrid models ([Attention Mechanisms](/tensorflow/advanced/attention-mechanisms)).

Challenges and Considerations

CapsNets face challenges like:

  • Computational Cost: Dynamic routing is resource-intensive. Optimize with [XLA Acceleration](/tensorflow/fundamentals/xla-acceleration).
  • Scalability: Scaling to large images requires efficient data pipelines ([Data Pipeline Scaling](/tensorflow/intermediate/data-pipeline-scaling)).
  • Interpretability: Use tools like the What-If Tool to understand predictions ([What-If Tool](/tensorflow/production/what-if-tool)).

Stay updated with TensorFlow’s advancements (TensorFlow Roadmap).

Conclusion

Capsule Networks offer a promising alternative to CNNs, capturing spatial hierarchies and improving generalization. By implementing CapsNets in TensorFlow, you can tackle complex vision tasks with robust models. Experiment with different datasets, architectures, and optimization strategies to fully leverage their potential.