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
:
from tensorflow.keras import layersimport tensorflow_probability as tfpclass 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_unitsself.number_visible_units = number_visible_unitsself.learning_rate = learning_rateself.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:
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
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:
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.