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

Answer 1

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