Error when not using the defined layer in call method in Keras model with WandB integration

When defining a new Keras model and integrating it with WandB (Weights and Biases), I encountered an error related to a specific layer. The error message I received is as follows:

File "/miniconda3/envs/tf/lib/python3.9/site-packages/wandb/data_types.py", line 1504, in from_keras
    node = Node.from_keras(layers[i])
  File "/miniconda3/envs/tf/lib/python3.9/site-packages/wandb/data_types.py", line 1731, in from_keras
    node.num_parameters = layer.count_params()
ValueError: You tried to call `count_params` on layer conv1d, but the layer isn't built. You can build it manually via: `conv1d.build(batch_input_shape)`.

Steps to Reproduce:

Define a new Keras model.
Include a specific layer within the __init__ method of the model.
Do not use the aforementioned layer in the call method of the model.
Integrate the Keras model with WandB for tracking.

Expected Behavior: The model should integrate with WandB without any errors, even if there are layers defined in the __init__ method that is not used in the call method.

Actual Behavior: The error mentioned above is raised when trying to integrate the Keras model with WandB.

Code Snippet:

import tensorflow as tf
import wandb

class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1d = tf.keras.layers.Conv1D(64, kernel_size=3, activation='relu')
        # Other layers...

    def call(self, inputs):
        # Only using other layers, not `conv1d`
        return inputs

# WandB initialization
wandb.init(project="my-project", mode='disabled')
config = wandb.config

# Create and integrate the model
model = MyModel()

wandb_callback = wandb.keras.WandbCallback(
    monitor="val_loss", 
    verbose=0, 
    mode="min", 
    save_model=False
)

x_train, y_train = tf.random.uniform([100, 10]), tf.random.uniform([100, 10])
wandb.config.update(config)  # Update config with any hyperparameters
model.compile(optimizer='adam', loss='mse')
model.fit(x_train, y_train, epochs=5, callbacks=[wandb_callback])
wandb.log({"example_metric": 0.85})
wandb.finish()

Please let me know if you need any further information or clarification.

Environment

WandB version: 0.15.8

OS: Ubuntu 22.04

Python version: 3.9.16

TensorFlow Version: [2.12.0]

Solution

The error is not specific to wandb, the problem is in the users' code since the model and inner layers does not have a specified input shape. To reproduce the error without wandb:

import tensorflow as tf

class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1d = tf.keras.layers.Conv1D(64, kernel_size=3, activation='relu')
        # Other layers...

    def call(self, inputs):
        # Only using other layers, not `conv1d`
        return inputs

# Create and integrate the model
model = MyModel()
    
model.conv1d.count_params()

This will output: ValueError: You tried to call count_paramson layer conv1d, but the layer isn't built. You can build it manually via:conv1d.build(batch_input_shape).

To fix it, just call build on your model with the shape of the input data minus the batch shape, for example here I use (32, 32, 3):

import tensorflow as tf

class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1d = tf.keras.layers.Conv1D(64, kernel_size=3, activation='relu')
        # Other layers...

    def call(self, inputs):
        # Only using other layers, not `conv1d`
        return inputs

# Create and integrate the model
model = MyModel()

model.build((32,32,3))

model.summary()

The error is gone and you can use the model. In any case you have not built the inner layers, which you should also do if you plan on using them somehow, just by calling build on them:

import tensorflow as tf

class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1d = tf.keras.layers.Conv1D(64, kernel_size=3, activation='relu')
        # Other layers...
        self.conv1d.build((32,32,3))

    def call(self, inputs):
        # Only using other layers, not `conv1d`
        return inputs

# Create and integrate the model
model = MyModel()

model.build((32,32,3))

print(model.conv1d.count_params())

This correctly prints 640 for the given input shape. In the end I reiterate, the problem is not in wandb but in your model as it has unused layers that do not have defined input shapes, just calling build on them will make everything work.