Unexpected Behavior when Replacing Layers in CustomModel built with Subclass API

[ariG23498]

I have created a CustomModel using the subclass API

class CustomModel(keras.Model):
    def __init__(self, num_classes=10):
        super().__init__()
        self.num_classes = num_classes
        self.stem = keras.layers.Conv2D(32, 3, strides=2, padding='same', name="stem")
        self.head = keras.layers.Dense(num_classes, name="head")
    
    def call(self, inputs):
        x = self.stem(inputs)
        x = keras.layers.Flatten()(x)
        x = self.head(x)
        return x

I wanted to replace (read swap) a layer from the custom model with a new layer.

model = CustomModel()
model.head = keras.layers.Dense(100, name="head")

Now when I hit model.summary I see both the layers added to the Model.

model.summary()

Model: "custom_model_3"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                         ┃ Output Shape                ┃         Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
│ stem (Conv2D)                        │ ?                           │     0 (unbuilt) │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ head (Dense)                         │ ?                           │     0 (unbuilt) │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ head (Dense)                         │ ?                           │     0 (unbuilt) │
└──────────────────────────────────────┴─────────────────────────────┴─────────────────┘
 Total params: 0 (0.00 B)
 Trainable params: 0 (0.00 B)
 Non-trainable params: 0 (0.00 B)

When I take a similar approach with torch.nn.Module, the layers get swapped out.

class CustomModel(torch.nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.num_classes = num_classes
        self.stem = torch.nn.Conv2d(32, 32, 3, stride=2)
        self.head = torch.nn.Linear(32, num_classes)
    
    def forward(self, inputs):
        x = self.stem(inputs)
        x = keras.layers.Flatten()(x)
        x = self.head(x)
        return x

model = CustomModel()
model.head = torch.nn.Linear(32, 100)
print(model)

CustomModel(
  (stem): Conv2d(32, 32, kernel_size=(3, 3), stride=(2, 2))
  (head): Linear(in_features=32, out_features=100, bias=True)
)

I know that I can use the setattr method to set the keras.Model attribute. But here I wanted to check whether this is a supposed workflow.

Any help on this would be grately appreciated.

My opinion: We should warn the user if they want to set the attribute of a model (mainly layers). I feel adding a layer to the model (which is not intuitive at all) introduces a silent bug.

[SuryanarayanaY]

Hi @ariG23498 ,

Keras won't allow to adding new layers once model.build() is called. Before model build call it is accepting the duplicate layers though. This behaviour is different wrt pytorch which swaps the duplicate layers instead of adding as new.Attached gist for reference.

[ariG23498]

Hi @ariG23498 ,

Keras won't allow to adding new layers once model.build() is called. Before model build call it is accepting the duplicate layers though. This behaviour is different wrt pytorch which swaps the duplicate layers instead of adding as new.Attached gist for reference.

Hey @SuryanarayanaY I did look into the GIST. I did not really understand the reason of the following code with the torch backend

model.build(input_shape=(12,4))
model.head_2 = keras.layers.Dense(10, name="head_2")
model.summary()

AFAIK the CustomModel is a torch.nn.Module. This will not include the build() method in the first place.

A better side by side analysis of the situation would be where we run the same code as follows, with different backends.

class CustomModel(keras.Model):
    def __init__(self, num_classes=10):
        super().__init__()
        self.num_classes = num_classes
        self.stem = keras.layers.Conv2D(32, 3, strides=2, padding='same', name="stem")
        self.head = keras.layers.Dense(num_classes, name="head")

    def call(self, inputs):
        x = self.stem(inputs)
        x = keras.layers.Flatten()(x)
        x = self.head(x)
        return x

model = CustomModel()
model.head = keras.layers.Dense(100, name="head")
print(model.summary())

model.build(input_shape=(12, 4))
model.head = keras.layers.Dense(100, name="head")
print(model.summary())

I did run this experiment and I saw that all the backends have the same behaviour.

[LarsKue]

Bumping this issue as it appears unresolved. Can we expect a fix soon?

Hoping that it helps, I have investigated the issue even further and the problem is that the layers are added to the _layers attribute when assigned, but not removed when overwritten.

class Layer(keras.Layer):
    def __init__(self):
        super().__init__()

        self.layer = keras.layers.Dense(2)
        self.layer = keras.layers.Dense(2)


layer = Layer()

print(layer._layers)
>>> [<Dense name=dense, built=False>, <Dense name=dense_1, built=False>]

This can lead to very confusing error messages when using a custom train_step. Note that this also happens under jax currently, I am simply using torch as an example:

import os
os.environ["KERAS_BACKEND"] = "torch"

import keras
import torch


class Model(keras.Model):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def build(self, input_shape):
        self.layer.build(input_shape)

    def compute_loss(self, data):
        return self.layer.compute_loss(data)

    def train_step(self, data):
        with torch.enable_grad():
            loss = self.compute_loss(data)

        # noinspection PyUnresolvedReferences
        self.zero_grad()
        loss.backward()

        trainable_weights = self.trainable_weights[:]
        gradients = [v.value.grad for v in trainable_weights]

        # Update weights
        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        self._loss_tracker.update_state(loss)

        return {"loss": loss}


class Layer(keras.Layer):
    def __init__(self):
        super().__init__()

        self.layer = keras.layers.Dense(2)
        self.layer = keras.layers.Dense(2)

    def call(self, inputs):
        return self.layer(inputs)

    def build(self, input_shape):
        self.layer.build((2,))

    def compute_loss(self, data):
        x, y = data
        return keras.losses.mean_squared_error(self.layer(x), y)


model = Model(Layer())

x = keras.ops.zeros((32, 2))
y = keras.ops.zeros((32, 2))

model.compile(optimizer="adam")
model.fit(x, y)

yields

RuntimeError: Unable to automatically build the model. Please build it yourself before calling fit/evaluate/predict. A model is 'built' when its variables have been created and its `self.built` attribute is True. Usually, calling the model on a batch of data is the right way to build it.
Exception encountered:
'Exception encountered when calling Model.call().

Model Model does not have a `call()` method implemented.

Arguments received by Model.call():
  • args=('torch.Tensor(shape=torch.Size([32, 2]), dtype=float32)',)
  • kwargs=<class 'inspect._empty'>'

In this case, the issue is that the Model is considered as not built because the layers that were overwritten did not get built in the Layer's build method.

This is especially problematic when wrapping layers conditionally. Here is a real use case:

@serializable(package="bayesflow.networks")
class SkipRecurrentNet(keras.Model):
    def __init__(
        self,
        ...
        bidirectional: bool = True,
        ...
    ):
        ...

        self.recurrent = ...  # some keras.Layer
        self.skip_recurrent = ...  # some keras.Layer
        if bidirectional:
            self.recurrent = keras.layers.Bidirectional(self.recurrent)
            self.skip_recurrent = keras.layers.Bidirectional(self.skip_recurrent)

The current work-around is to only assign the attribute to self after it is fully constructed, i.e. use local variables and assign to self only after the if statement.

However, considering how unintuitive the error is and especially its inconsistency across backends, I would expect that this should be fixed in keras, e.g. by tracking a dictionary of assigned layers rather than simply a list, or by constructing the _layers list on the fly once it is accessed.