Trainable weights in custom layers?

Question

I was learning making custom layers in tensor flow but could not find out how to add trainable weights for example

class Linear(layers.Layer):

  def __init__(self, units = 32, **kwargs):
    super().__init__(kwargs)
    self.units = units

  def build(self, input_shape):
    self.layer = layers.Dense(self.units, trainable= True)
    super().build(input_shape)  
  
  def call(self, inputs):
    return self.layer(inputs)

Now if I do

linear_layer = Linear(8)

x = tf.ones(shape =(4,3))
y = linear_layer(x)

print(linear_layer.trainable_variables)

I get an empty matrix and thus during gradient calculation I get no gradients, my question is how to create custom layers in a way that default keras layers are also trainable in that. One more thing if I do linear_layer.weights then it give me the weights, it means there is some problem with trainable weights. My mind is stuck on that

If the code below solved your problem. please mark it as the correct answer. — Tou You, Jul 25 '21 at 09:34

Tou You · Accepted Answer · 2021-07-24T17:34:00.037

To get trainable variables you have to access the "layer" attribute of your custom layer:

linear_layer = Linear(8)

x = tf.ones(shape =(4,3))
y = linear_layer(x)

print(linear_layer.layer.trainable_variables)

note that you just create a pre_built layer (Dense) in the build method instead of create the weights of your custom layer. look at link https://www.tensorflow.org/tutorials/customization/custom_layers

Trainable weights in custom layers?

1 Answers1