[QUESTION] Chapter 12, How to pass training=True argument used in call() of custom layer

[michaelmiscanuk]

Hi, on page 424 you show that we can control the layer based on argument training set to True or False - but I dont know how to set it to True (or how to pass a value True to the parameter training in call() moethod ??

I tries passing it to the layer, to the compile and to the fit - but I got errors.

Here is my example code, which I run and see that the default is used, because output of the layer is not changed, this can be seen in output of the custom model, where I input 3, and output is same - 3.

# %%
import tensorflow as tf
from tensorflow import keras
import numpy as np

### DATA
# x values
x = tf.constant([[1.0], [2.0], [3.0]])

# y values
y = tf.constant([[2.0], [6.0], [4.0]])

### PLOT
# Plot all the data
import matplotlib.pyplot as plt

plt.scatter(x, y, c="b")
plt.show()

# %%
# CUSTOM LAYER
class AddGaussianNoise(keras.layers.Layer):
    def __init__(self, stddev, **kwargs):
        super().__init__(**kwargs)
        self.stddev = stddev

    def call(self, X, training=None):
        if training:
            # noise = tf.random.normal(tf.shape(X), stddev=self.stddev)
            noise = 5
            return X + noise
        else:
            return X

    def compute_output_shape(self, batch_input_shape):
        return batch_input_shape


class CustomCallback(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(model.layers[0].output)


tf.random.set_seed(42)
np.random.seed(42)

training_bool = True

model = keras.models.Sequential(
    [
        AddGaussianNoise(50, name="add_5"),
        keras.layers.Dense(1),
    ]
)

lr0 = 0.1
optimizer = keras.optimizers.Nadam(learning_rate=lr0)
model.compile(
    loss="mse",
    optimizer=optimizer,
    metrics=["mae"],
)

history = model.fit(
    x,
    y,
    epochs=1,
    validation_data=(x, y),
    callbacks=[CustomCallback()],
)

# Visualize how the trained model performs
plt.scatter(x, y, c="b")
plt.scatter(x, model.predict(x), c="r")
plt.plot(x, model.predict(x), c="r")
plt.show()


### Model to Show Output of Some Layer
from keras.models import Model

layer_name = "add_5"
intermediate_layer_model = Model(
    inputs=model.input,
    outputs=model.get_layer(layer_name).output,
)
intermediate_output = intermediate_layer_model.predict([3])
intermediate_output
# array([[3.]], dtype=float32)

Can you help? Thank you.

[ageron]

Hi @michaelmiscanuk ,

The training argument is automatically set to True by Keras during training, you don't need to do anything. And it's automatically left to the default value otherwise.

If you want to make a prediction but forcing training=True, you can call the layer as if it were a function, and pass the training argument:

gaussian_noise_layer = model.layers[0]
assert gaussian_noise_layer(tf.constant([[1.0]])) == 1.0
assert gaussian_noise_layer(tf.constant([[1.]]), training=True) == 6.0

Hope this helps.

[ageron]

Note: if you want to be sure that training is set to True, you can use tf.print() to print a log message, for example:

class AddGaussianNoise(keras.layers.Layer):
    def __init__(self, stddev, **kwargs):
        super().__init__(**kwargs)
        self.stddev = stddev

    def call(self, X, training=None):
        if training:
            # noise = tf.random.normal(tf.shape(X), stddev=self.stddev)
            noise = 5
            tf.print("TRAINING!")  # <= LOG MESSAGE
            return X + noise
        else:
            return X

    def compute_output_shape(self, batch_input_shape):
        return batch_input_shape

If you train a model containing this layer, you will see the message "TRAINING!" displayed at every batch.

Side note: TensorFlow is generally able to determine the shape of the output automatically now, so the compute_output_shape method is not required.