Custom layers with default layers composition hide their weights

[shkarupa-alex]

Here is an example with Keras 2 vs 3 https://colab.research.google.com/drive/1WKUPCbG14Hdn39pQLigfW81-0Asw9Shr?usp=sharing

As result model pretrained weights can't be loaded.

[mehtamansi29]

Hi @shkarupa-alex -

You can create custom layers like this, it will load pretrained weights in keras3.

from keras import layers, models

class MyCompostionLayer(layers.Layer):
  def __init__(self, input_shape):
    super(MyCompostionLayer, self).__init__()
    self.fc1 = layers.Dense(int(input_shape[-1] * 4), name='fc1')
    self.fc2 = layers.Dense(input_shape[-1], name='fc2')

  def call(self, inputs):
      outputs = self.fc1(inputs)
      outputs = self.fc2(outputs)
      return outputs

  def compute_output_shape(self, input_shape):
    return input_shape

x = layers.Input(shape=(None, 4), dtype='float32')
y = MyCompostionLayer(x.shape)(x)

m = models.Model(x, y)
print('Number of weights:', len(m.weights))
print(m.weights)
m.summary()

Attached gist for your reference.

[shkarupa-alex]

Hi @shkarupa-alex -

You can create custom layers like this, it will load pretrained weights in keras3.

from keras import layers, models

class MyCompostionLayer(layers.Layer):
  def __init__(self, input_shape):
    super(MyCompostionLayer, self).__init__()
    self.fc1 = layers.Dense(int(input_shape[-1] * 4), name='fc1')
    self.fc2 = layers.Dense(input_shape[-1], name='fc2')

  def call(self, inputs):
      outputs = self.fc1(inputs)
      outputs = self.fc2(outputs)
      return outputs

  def compute_output_shape(self, input_shape):
    return input_shape

x = layers.Input(shape=(None, 4), dtype='float32')
y = MyCompostionLayer(x.shape)(x)

m = models.Model(x, y)
print('Number of weights:', len(m.weights))
print(m.weights)
m.summary()

Attached gist for your reference.

This is not a solution for the case when custom layer requires some information about inputs (like number of channels in input_shape)

[mehtamansi29]

Hi @shkarupa-alex -

If want to get some information about inputs (like number of channels in input_shape) in custom layer like this:

class MyCompostionLayer(layers.Layer):
  def __init__(self, input_shape):
    super(MyCompostionLayer, self).__init__()
    print(f'Input Shape:{input_shape[1:]}')       # Input Shape:(None, 4)
    print(f'No. of channels:{input_shape[-1]}')      #No. of channels:4
    self.fc1 = layers.Dense(int(input_shape[-1] * 4), name='fc1')
    self.fc2 = layers.Dense(input_shape[-1], name='fc2')

  def call(self, inputs):
      outputs = self.fc1(inputs)
      outputs = self.fc2(outputs)
      return outputs

  def compute_output_shape(self, input_shape):
    return input_shape

x = layers.Input(shape=(None, 4), dtype='float32')
y = MyCompostionLayer(x.shape)(x)

m = models.Model(x, y)
print('Number of weights:', len(m.weights))
print(m.weights)
print(f'Input Shape:{x.shape[1:]}')    #Input Shape:(None, 4)
print(f'No. of channels:{x.shape[-1]}')     #No. of channels:4
m.summary()

Please let me know if you need more help. Thanks..!!

[shkarupa-alex]

Hi @shkarupa-alex -

If want to get some information about inputs (like number of channels in input_shape) in custom layer like this:

class MyCompostionLayer(layers.Layer):
  def __init__(self, input_shape):
    super(MyCompostionLayer, self).__init__()
    print(f'Input Shape:{input_shape[1:]}')       # Input Shape:(None, 4)
    print(f'No. of channels:{input_shape[-1]}')      #No. of channels:4
    self.fc1 = layers.Dense(int(input_shape[-1] * 4), name='fc1')
    self.fc2 = layers.Dense(input_shape[-1], name='fc2')

  def call(self, inputs):
      outputs = self.fc1(inputs)
      outputs = self.fc2(outputs)
      return outputs

  def compute_output_shape(self, input_shape):
    return input_shape

x = layers.Input(shape=(None, 4), dtype='float32')
y = MyCompostionLayer(x.shape)(x)

m = models.Model(x, y)
print('Number of weights:', len(m.weights))
print(m.weights)
print(f'Input Shape:{x.shape[1:]}')    #Input Shape:(None, 4)
print(f'No. of channels:{x.shape[-1]}')     #No. of channels:4
m.summary()

Please let me know if you need more help. Thanks..!!

Yes, i need more help. If i would like to hardcode every shape like in PyTorch, i'd use it directly without Keras.

But i want to write reusable layers. Just like in Keras v2 where provided example works as expected.

Passing shape manually in constructor is not a solution.

[james77777778]

Hey @shkarupa-alex

This should meet your needs:

from keras import layers
from keras import models

class MyCompostionLayer(layers.Layer):
    def build(self, input_shape):
        self.fc1 = layers.Dense(int(input_shape[-1] * 4), name="fc1")
        self.fc2 = layers.Dense(input_shape[-1], name="fc1")

        # Manually build sublayers
        self.fc1.build(input_shape)
        input_shape = self.fc1.compute_output_shape(input_shape)
        self.fc2.build(input_shape)

    def call(self, inputs):
        outputs = self.fc1(inputs)
        outputs = self.fc2(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape

x = layers.Input(shape=(None, 4), dtype="float32")
y = MyCompostionLayer()(x)
m = models.Model(x, y)

print("Number of weights", len(m.weights), "\n")
m.summary()

The idea is that we need to manually build the sublayers if we instantiate them in build. Just like: https://github.com/keras-team/keras/blob/dfc4f788720ee58322c927ef061501b48a47eac3/keras/src/layers/attention/multi_head_attention.py#L234-L243

Another approach is to feed the inputs to build them:

from keras import layers
from keras import models
from keras import ops

class MyCompostionLayer(layers.Layer):
    def build(self, input_shape):
        self.fc1 = layers.Dense(int(input_shape[-1] * 4), name="fc1")
        self.fc2 = layers.Dense(input_shape[-1], name="fc1")

    def call(self, inputs):
        outputs = self.fc1(inputs)
        outputs = self.fc2(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape

x = layers.Input(shape=(None, 4), dtype="float32")
y = MyCompostionLayer()(x)
m = models.Model(x, y)

m(ops.ones((1, 1, 4)))  # Eager call to build sublayers

print("Number of weights", len(m.weights), "\n")
m.summary()

[shkarupa-alex]

Resolved within TestCase.run_layer_test

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No