Making some of the kernel parameters of the dense layer non-trainable

Viewed 270

Consider a simple model as follows:

input = keras.layers.Input(shape=(1,))
l1 = tf.keras.layers.Dense(units=4,activation='relu',use_bias=False)(input)
out = tf.keras.layers.Dense(units=4,activation=None,use_bias=False)(l1)
model = tf.keras.Model(inputs=input, outputs=out)

So for the last layer, the kernel has a shape of (4,4). Although for my implementation, I would only want to keep some of the parameters in this kernel to be trainable.

For instance, in usual cases, my kernel for the last layer looks like this:

enter image description here

But I would like to have something like this:

enter image description here

I want to have only non-zero entries as trainable. Can someone please help with this?

1 Answers

The simplest way to do this is using kernel_initializer='zeros' with a custom kernel_constraint in the last Dense.

Here is my proposal for the custom kernel_constraint. my_constraints is a matrix of 0/1 built outside that must be passed as argument.

class MyConstraint(tf.keras.constraints.Constraint):

  def __init__(self, my_constraints):
    self.my_constraints = my_constraints

  def __call__(self, w):
    assert tf.shape(w).shape == tf.shape(self.my_constraints).shape
    return w * tf.cast(self.my_constraints, dtype=tf.float32)

  def get_config(self):
    return {'my_constraints': self.my_constraints}

Below is how to build and run the model with some dummy data:

my_constraints = tf.constant([[1,0,1,0],[0,1,0,1],[1,0,1,0],[0,1,0,1]], dtype=tf.float32)
X = np.random.uniform(-10,10, (100,1))
y = np.random.uniform(-10,10, (100,4))

input = Input(shape=(1,))
l1 = Dense(units=4, activation='relu', name='layer1', use_bias=False)(input)
out = Dense(units=4, activation=None , name='layer2', use_bias=False, 
            kernel_constraint=MyConstraint(my_constraints), kernel_initializer='zeros')(l1)
model = Model(inputs=input, outputs=out)
model.compile('adam','mse')

model.fit(X,y, epochs=100)

To inspect the results:

>>> print(model.get_weights()[-1] == 0)
[[False  True False  True]
 [ True False  True False]
 [False  True False  True]
 [ True False  True False]]

>>> print(model.get_weights()[-1])
[[-0.00230937 -0.         -0.01893261  0.        ]
 [-0.         -0.07535305 -0.         -0.01603868]
 [-0.16068968  0.          0.23261748 -0.        ]
 [-0.         -0.07739028 -0.         -0.01623872]]
Related