- Published on

# A Common Pitfall When Writing Subclass Models in Keras

- Authors
- Name
- Luke Wood
- @luke_wood_ml

## Overview

The other day while benchmarking the KerasCV RetinaNet I was shocked to see that the prediction and regression heads were both not training. The model achieved mediocre scores; but the weights were not being updated. After almost a week of debugging I found the culprit.

TLDR: subclass models cannot track the weights used in closure-style components; which are a commonly used abstraction in the `keras_cv.models`

subpackage.

## Demonstration

Lets demonstrate this edge case:

```
import tensorflow as tf
from tensorflow import keras
import numpy as np
class CustomModel(keras.Model):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.dense_1 = keras.layers.Dense(10)
self.dense_2 = keras.layers.Dense(1)
def call(self, x, training=None):
x = self.dense_1(x)
return self.dense_2(x)
def train_step(self, data):
# Unpack the data. Its structure depends on your model and
# on what you pass to `fit()`.
x, y = data
with tf.GradientTape() as tape:
y_pred = self(x, training=True) # Forward pass
# Compute the loss value
# (the loss function is configured in `compile()`)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
return {m.name: m.result() for m in self.metrics}
```

This is a properly written subclass model. You can fit it using:

```
model = CustomModel()
model.compile(optimizer="adam", loss="mse")
# Just use `fit` as usual -- you can use callbacks, etc.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)
```

```
Epoch 1/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1832
Epoch 2/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1476
Epoch 3/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1306
Epoch 4/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1150
Epoch 5/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1073
<keras.callbacks.History at 0x12e7dc670>
```

Cool, it works. Thats all there is to it. Lets re-write this using the closure pattern to abstract the dense layers into a single block.

```
def DenseBlock():
dense_1 = keras.layers.Dense(10)
dense_2 = keras.layers.Dense(1)
def apply(x, training=None):
x = dense_1(x)
return dense_2(x)
return apply
```

Great, and lets use it in our model:

```
class CustomModel(keras.Model):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.dense_block = DenseBlock()
def call(self, x, training=None):
return self.dense_block(x)
def train_step(self, data):
# Unpack the data. Its structure depends on your model and
# on what you pass to `fit()`.
x, y = data
with tf.GradientTape() as tape:
y_pred = self(x, training=True)
# Compute the loss value
# (the loss function is configured in `compile()`)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars) # Broken
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
return {m.name: m.result() for m in self.metrics}
model = CustomModel()
model.compile(optimizer="adam", loss="mse")
# Just use `fit` as usual -- you can use callbacks, etc.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)
```

```
Epoch 1/5
32/32 [==============================] - 0s 974us/step - loss: 0.5985
Epoch 2/5
32/32 [==============================] - 0s 1ms/step - loss: 0.5985
Epoch 3/5
32/32 [==============================] - 0s 983us/step - loss: 0.5985
Epoch 4/5
32/32 [==============================] - 0s 1ms/step - loss: 0.5985
Epoch 5/5
32/32 [==============================] - 0s 1ms/step - loss: 0.5985
<keras.callbacks.History at 0x12eee7cd0>
```

Notice something? The model doesn't learn.

Why? Lets dig deeper:

```
print(model.trainable_variables)
```

```
[]
```

There lies the culprit: the model doesn't know about the variables in `DenseBlock()`

. This is because Keras cannot check the internals of the assigned function definition and recognize that it represents a layer.

To fix this just construct a custom layer, functional model, or sequential represnting the same structure:

```
def DenseBlock():
return keras.Sequential([Dense(10), Dense(1)])
```

Or:

```
class DenseBlock(keras.layers.Layer):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.dense_1 = keras.layers.Dense(10)
self.dense_2 = keras.layers.Dense(1)
def call(self, x, training=None):
x = self.dense_1(x)
return self.dense_2(x)
```

Thats all, now our custom model will work again:

```
class CustomModel(keras.Model):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.dense_block = DenseBlock()
def call(self, x, training=None):
return self.dense_block(x)
def train_step(self, data):
# Unpack the data. Its structure depends on your model and
# on what you pass to `fit()`.
x, y = data
with tf.GradientTape() as tape:
y_pred = self(x, training=True) # Forward pass
# Compute the loss value
# (the loss function is configured in `compile()`)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
return {m.name: m.result() for m in self.metrics}
model = CustomModel()
model.compile(optimizer="adam", loss="mse")
# Just use `fit` as usual -- you can use callbacks, etc.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)
```

```
Epoch 1/5
32/32 [==============================] - 0s 1ms/step - loss: 0.3773
Epoch 2/5
32/32 [==============================] - 0s 1ms/step - loss: 0.2919
Epoch 3/5
32/32 [==============================] - 0s 1ms/step - loss: 0.2321
Epoch 4/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1915
Epoch 5/5
32/32 [==============================] - 0s 1ms/step - loss: 0.1594
<keras.callbacks.History at 0x12ef8c3d0>
```

## Conclusion

The takeaway is that subclass models can't track variables used in closure style functional layers.

Unless you really know Keras inside and out, just use `tf.keras.layers.Layer`

when building a sub-class model. Keep closure style functions to the functional model building API.