title trainer pattern nkovela1https nkovela1 github io date created 20220919 last modified 20220926 description guide on how to share a custom training step across multiple keras models accelerator gpu introduction this example shows how to create a custom training step using the trainer pattern which can then be shared across multiple keras models this pattern overrides the trainstep method of the keras model class allowing for training loops beyond plain supervised learning the trainer pattern can also easily be adapted to more complex models with larger custom training steps such as this endtoend gan modelhttps keras ioguidescustomizingwhathappensinfitwrappingupanendtoendganexample by putting the custom training step in the trainer class definition setup load mnist dataset and standardize the data define the trainer class a custom training and evaluation step can be overriding the trainstep and teststep method of a model subclass create loss and metrics here list metrics here compute loss value compute gradients update weights update metrics return a dict mapping metric names to current value inference step update metrics equivalent to call of the wrapped keras model define multiple models to share the custom training step let s define two different models that can share our trainer class and its custom trainstep a model defined using sequential api a model defined using functional api create trainer class objects from the models compile and fit the models to the mnist dataset introduction this example shows how to create a custom training step using the trainer pattern which can then be shared across multiple keras models this pattern overrides the train_step method of the keras model class allowing for training loops beyond plain supervised learning the trainer pattern can also easily be adapted to more complex models with larger custom training steps such as this end to end gan model https keras io guides customizing_what_happens_in_fit wrapping up an endtoend gan example by putting the custom training step in the trainer class definition setup load mnist dataset and standardize the data define the trainer class a custom training and evaluation step can be overriding the train_step and test_step method of a model subclass create loss and metrics here list metrics here forward pass compute loss value compute gradients update weights update metrics return a dict mapping metric names to current value inference step update metrics equivalent to call of the wrapped keras model define multiple models to share the custom training step let s define two different models that can share our trainer class and its custom train_step a model defined using sequential api a model defined using functional api create trainer class objects from the models compile and fit the models to the mnist dataset
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras


mnist = keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0





class MyTrainer(keras.Model):
    def __init__(self, model):
        super().__init__()
        self.model = model
        
        self.loss_fn = keras.losses.SparseCategoricalCrossentropy()
        self.accuracy_metric = keras.metrics.SparseCategoricalAccuracy()

    @property
    def metrics(self):
        
        return [self.accuracy_metric]

    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self.model(x, training=True)  
            
            loss = self.loss_fn(y, y_pred)

        
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        
        for metric in self.metrics:
            metric.update_state(y, y_pred)

        
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        x, y = data

        
        y_pred = self.model(x, training=False)

        
        for metric in self.metrics:
            metric.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

    def call(self, x):
        
        x = self.model(x)
        return x





model_a = keras.models.Sequential(
    [
        keras.layers.Flatten(input_shape=(28, 28)),
        keras.layers.Dense(256, activation="relu"),
        keras.layers.Dropout(0.2),
        keras.layers.Dense(10, activation="softmax"),
    ]
)


func_input = keras.Input(shape=(28, 28, 1))
x = keras.layers.Flatten(input_shape=(28, 28))(func_input)
x = keras.layers.Dense(512, activation="relu")(x)
x = keras.layers.Dropout(0.4)(x)
func_output = keras.layers.Dense(10, activation="softmax")(x)

model_b = keras.Model(func_input, func_output)



trainer_1 = MyTrainer(model_a)
trainer_2 = MyTrainer(model_b)



trainer_1.compile(optimizer=keras.optimizers.SGD())
trainer_1.fit(
    x_train, y_train, epochs=5, batch_size=64, validation_data=(x_test, y_test)
)

trainer_2.compile(optimizer=keras.optimizers.Adam())
trainer_2.fit(
    x_train, y_train, epochs=5, batch_size=64, validation_data=(x_test, y_test)
)