Search⌘ K
Home
Courses
Generative AI with Python and TensorFlow 2

Creating an RBM using the TensorFlow Keras layers API

Explore how to implement a Restricted Boltzmann Machine (RBM) with TensorFlow 2's Keras layers API. Understand constructing custom layers, applying contrastive divergence, and training the RBM to generate digit images from the MNIST dataset. Gain hands-on experience leveraging TensorFlow's gradient tape and efficient sampling methods to optimize model training and visualize learned features.

We'll cover the following...

Now that we have an appreciation of some of the theoretical underpinnings of the RBM, let's learn how we can implement it using the TensorFlow 2 library. For this purpose, we’ll represent the RBM as a custom layer type using the Keras layers API. We’ll implement the whole model in code, showing how we can leverage TensorFlow 2’s gradient tape functionality to implement CD as a custom learning algorithm.

Implementing the RBM

Firstly, we extend tf.keras.layer:

Python 3.10.4
from tensorflow.keras import layers
import tensorflow_probability as tfp
class RBM(layers.Layer):
    def __init__(self, number_hidden_units=10, number_visible_units=None, 
                learning_rate=0.1, cd_steps=1):
        
        super().__init__()
        self.number_hidden_units = number_hidden_units
        self.number_visible_units = number_visible_units
        self.learning_rate = learning_rate
        self.cd_steps = cd_steps

We input a number of hidden units, visible units, a learning rate for CD updates, and the number of steps to take with each CD pass. For the layers API, we are only required to implement two functions: build() and call(). build() is executed when we call model.compile(), and is used to initialize the weights of the network, including inferring the right size of the weights given the input dimensions:

Python 3.10.4
def build(self, input_shape):
    if not self.number_visible_units:
        self.number_visible_units = input_shape[-1]
        
        self.w_rec = self.add_weight(shape=(self.number_visible_units,
                          self.number_hidden_units),
                          initializer='random_normal',
                          trainable=True)
        self.w_gen = self.add_weight(shape=(self.number_hidden_units,
                           self.number_visible_units),
                           initializer='random_normal',
                           trainable=True)
        self.hb = self.add_weight(shape=(self.number_hidden_units, ),
                           initializer='random_normal',
                           trainable=True)
        self.vb = self.add_weight(shape=(self.number_visible_units, ),
                           initializer='random_normal',
                           trainable=True)

We also need a way to perform both forward and reverse samples from the model. For the forward pass, we need to compute sigmoidal activations from the input, and then stochastically turn the hidden units on or off based on the activation probability between 11 and 00 given by that sigmoidal activation:

Python 3.10.4
def forward(self, x):
    return tf.sigmoid(tf.add(tf.matmul(x, self.w), self.hb))
def sample_h(self, x):
    u_sample = tfp.distributions.Uniform().sample((x.shape[1], self.hb.shape[-1]))
    return tf.cast((x) > u_sample, tf.float32)

Likewise, we need a way to sample in reverse for the visible units:

Python 3.10.4
def reverse(self, x):
    return tf.sigmoid(tf.add(tf.matmul(x, self.w_gen), self.vb))
def sample_v(self, x):
    u_sample = tfp.distributions.Uniform().sample((x.shape[1], self.vb.shape[-1]))
    return tf.cast(self.reverse(x) > u_sample, tf.float32)

We also implement call() in the RBM class, which provides the forward pass we would use if we were to use the fit() method of the ...