Creating an RBM using the TensorFlow Keras layers API

Learn how to implement Restricted Boltzmann Machines using TensorFlow 2.0.

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:

Press + to interact
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:

Press + to interact
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:

Press + to interact
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:

Press + to interact
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  ...

Get hands-on with 1400+ tech skills courses.