I am trying to do a simple custom training loop.
For some reason, the train_step() function gets ignored and a normal training loop is carried. I noticed when the word "Hello" is not printed when I run the script on my MacBook Pro with M1 chip. My old MacBook Pro (AMD) works perfectly and the code also worked perfectly on Google Colab.
The tensor flow version is 2.0.0 and Keras is 2.3.1. Thanks in advance for your help.
My Code is:
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import Model
# tf.config.run_functions_eagerly(True)
class CustomModel(Model):
# @tf.function
def train_step(self, data):
# Unpack the data. Its structure depends on your model and
# on what you pass to `fit()`.
x, y = data
print('Hello')
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)
# Compute gradients
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
# Update weights
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
# Update metrics (includes the metric that tracks the loss)
self.compiled_metrics.update_state(y, y_pred)
# Return a dict mapping metric names to current value
return {m.name: m.result() for m in self.metrics}
# Construct and compile an instance of CustomModel
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])
# Just use `fit` as usual
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=3)
Hello is not printed:
Hello is printed during training:
