[Learn about machine learning from the Keras] — 13. Custom initializers

This section will look at defining your own initializer. We can use a subClass to define an initializer, provided that the custom initializer class is callable, that is, the class must implement the call function. We can follow the example in the previous section and add a custom initializer to observe.

Examples are as follows:

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras.datasets import mnist
from tensorflow.keras.initializers import Initializer
import keras.optimizers as optimizers

(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
train_images = train_images.reshape((60000, 28 * 28))
train_images = train_images.astype("float32") / 255
test_images = test_images.reshape((10000, 28 * 28))
test_images = test_images.astype("float32") / 255

class ExampleRandomNormal(Initializer):

    def __init__(self, mean, stddev):
        self.mean = mean
        self.stddev = stddev
    
    def __call__(self, shape, dtype=None):
        return tf.random.normal(shape, mean=self.mean, stddev=self.stddev, dtype=dtype)

    def get_config(self): # To support serialization
        return {'mean': self.mean, 'stddev': self.stddev}

class SimpleDense(layers.Layer):

    def __init__(self, units=32):
        super(SimpleDense, self).__init__()
        self.units = units
        
    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                        initializer= ExampleRandomNormal( mean=0.0 , stddev=0.05 ),
                        trainable=True)
        self.b = self.add_weight(shape=(self.units,),
                                initializer='random_normal',
                                trainable=True)
    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b
        
model = Sequential([
SimpleDense(512),
layers.Dense(10, activation="softmax")
])


model.compile(optimizer= 
    optimizers.get( {"class_name": "rmsprop", "config": {"learning_rate" : 0.001} } ) ,
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"])
    
model.fit(train_images, train_labels, epochs=1, batch_size=128)

The picture above shows the operation of the initializer when the model is being built. The purple frame is the operation process of the custom initializer. The custom initializer class inherits the keras.initializers.Initializer class.

When a Layer using a custom initializer is being built, keras.engine.base_layer.Layer.add_weight will be called when setting the initializer, and then keras.initializers.get will be executed to set the initializer attribute of the Layer.

Next, when preparing to generate the weight tensor, call keras.engine.base_layer_utils.make_variable. For the use of initializer, it will continue to call ExampleRandomNormal.call. When keras.engine.base_layer_utils.make_variable generates a good weight tensor, it is passed back to the SimpleDense.build function for SimpleDense to set the weight attribute. Then continue to execute the build operation of the next connected Layer.

Therefore, it can be observed from the figure that the custom initializer must have its own __call__, so that keras.initializers.get can successfully detect that the initializer entity is callable, so that the build operation can continue. Then this __call__ function must be passed in shape, just like the __call__ function of the keras.initializers.Initializer parent class.

The above is the operation of custom initializer and things to pay attention to.