Are Keras custom layer parameters non-trainable by default?

user_14890777 0
Viewed 717

I built a simple custom layer in Keras and was surprised to find that the parameters were not set to trainable by default. I can get it to work by explicitly setting the trainable attribute. I can't explain why this is by looking at documentation or code. Is this how it is supposed to be or I am doing something wrong which is making the parameters non-trainable by default? Code:

import tensorflow as tf


class MyDense(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(MyDense, self).__init__(kwargs)
        self.dense = tf.keras.layers.Dense(2, tf.keras.activations.relu)

    def call(self, inputs, training=None):
        return self.dense(inputs)


inputs = tf.keras.Input(shape=10)
outputs = MyDense()(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name='test')
model.compile(loss=tf.keras.losses.MeanSquaredError())
model.summary()

Output:

Model: "test"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 10)]              0         
_________________________________________________________________
my_dense (MyDense)           (None, 2)                 22        
=================================================================
Total params: 22
Trainable params: 0
Non-trainable params: 22
_________________________________________________________________

If I change the custom layer creation like this:

outputs = MyDense(trainable=True)(inputs)

the output is what I expect (all parameters are trainable):

=================================================================
Total params: 22
Trainable params: 22
Non-trainable params: 0
_________________________________________________________________

then it works as expected and makes all the parameters trainable. I don't understand why that is needed though.

keras tensorflow python
1 Answers

No doubt, that's an interesting quirk.

When making a custom layer, a tf.Variable will be automatically included in the list of trainable_variable. You didn't use tf.Variable, but a tf.keras.layers.Dense object instead, which will not be treated as a tf.Variable, and not set trainable=True by default. However, the Dense object you used will be set to trainable. See:

MyDense().dense.trainable

True

If you used tf.Variable (as it should), it will be trainable by default.

import tensorflow as tf


class MyDense(tf.keras.layers.Layer):
    def __init__(self, units=2, input_dim=10):
        super(MyDense, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(input_dim, units), dtype="float32"),
            trainable=True,
        )
        b_init = tf.zeros_initializer()
        self.b = tf.Variable(
            initial_value=b_init(shape=(units,), dtype="float32"), trainable=True
        )

    def call(self, inputs, **kwargs):
        return tf.matmul(inputs, self.w) + self.b


inputs = tf.keras.Input(shape=10)
outputs = MyDense()(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name='test')
model.compile(loss=tf.keras.losses.MeanSquaredError())
model.summary()

Model: "test"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_11 (InputLayer)        [(None, 10)]              0         
_________________________________________________________________
my_dense_18 (MyDense)        (None, 2)                 22        
=================================================================
Total params: 22
Trainable params: 22
Non-trainable params: 0
_________________________________________________________________
user_109083750
Related