title customizing what happens in fit with pytorch fchollethttps twitter comfchollet date created 20230627 last modified 20230627 description overriding the training step of the model class with pytorch accelerator gpu introduction when you re doing supervised learning you can use fit and everything works smoothly when you need to take control of every little detail you can write your own training loop entirely from scratch but what if you need a custom training algorithm but you still want to benefit from the convenient features of fit such as callbacks builtin distribution support or step fusing a core principle of keras is progressive disclosure of complexity you should always be able to get into lowerlevel workflows in a gradual way you shouldn t fall off a cliff if the highlevel functionality doesn t exactly match your use case you should be able to gain more control over the small details while retaining a commensurate amount of highlevel convenience when you need to customize what fit does you should override the training step function of the model class this is the function that is called by fit for every batch of data you will then be able to call fit as usual and it will be running your own learning algorithm note that this pattern does not prevent you from building models with the functional api you can do this whether you re building sequential models functional api models or subclassed models let s see how that works setup this guide can only be run with the torch backend a first simple example let s start from a simple example we create a new class that subclasses keras model we just override the method trainstepself data we return a dictionary mapping metric names including the loss to their current value the input argument data is what gets passed to fit as training data if you pass numpy arrays by calling fitx y then data will be the tuple x y if you pass a torch utils data dataloader or a tf data dataset by calling fitdataset then data will be what gets yielded by dataset at each batch in the body of the trainstep method we implement a regular training update similar to what you are already familiar with importantly we compute the loss via self computeloss which wraps the losses functions that were passed to compile similarly we call metric updatestatey ypred on metrics from self metrics to update the state of the metrics that were passed in compile and we query results from self metrics at the end to retrieve their current value unpack the data its structure depends on your model and on what you pass to fit call torch nn module zerograd to clear the leftover gradients for the weights from the previous train step compute loss call torch tensor backward on the loss to compute gradients for the weights update weights update metrics includes the metric that tracks the loss return a dict mapping metric names to current value note that it will include the loss tracked in self metrics let s try this out construct and compile an instance of custommodel just use fit as usual going lowerlevel naturally you could just skip passing a loss function in compile and instead do everything manually in trainstep likewise for metrics here s a lowerlevel example that only uses compile to configure the optimizer we start by creating metric instances to track our loss and a mae score in init we implement a custom trainstep that updates the state of these metrics by calling updatestate on them then query them via result to return their current average value to be displayed by the progress bar and to be pass to any callback note that we would need to call resetstates on our metrics between each epoch otherwise calling result would return an average since the start of training whereas we usually work with perepoch averages thankfully the framework can do that for us just list any metric you want to reset in the metrics property of the model the model will call resetstates on any object listed here at the beginning of each fit epoch or at the beginning of a call to evaluate call torch nn module zerograd to clear the leftover gradients for the weights from the previous train step compute loss call torch tensor backward on the loss to compute gradients for the weights update weights compute our own metrics we list our metric objects here so that resetstates can be called automatically at the start of each epoch or at the start of evaluate construct an instance of custommodel we don t passs a loss or metrics here just use fit as usual you can use callbacks etc supporting sampleweight classweight you may have noticed that our first basic example didn t make any mention of sample weighting if you want to support the fit arguments sampleweight and classweight you d simply do the following unpack sampleweight from the data argument pass it to computeloss updatestate of course you could also just apply it manually if you don t rely on compile for losses metrics that s it unpack the data its structure depends on your model and on what you pass to fit call torch nn module zerograd to clear the leftover gradients for the weights from the previous train step compute loss call torch tensor backward on the loss to compute gradients for the weights update weights update metrics includes the metric that tracks the loss return a dict mapping metric names to current value note that it will include the loss tracked in self metrics construct and compile an instance of custommodel you can now use sampleweight argument providing your own evaluation step what if you want to do the same for calls to model evaluate then you would override teststep in exactly the same way here s what it looks like unpack the data compute predictions updates the metrics tracking the loss update the metrics return a dict mapping metric names to current value note that it will include the loss tracked in self metrics construct an instance of custommodel evaluate with our custom teststep wrapping up an endtoend gan example let s walk through an endtoend example that leverages everything you just learned let s consider a generator network meant to generate 28x28x1 images a discriminator network meant to classify 28x28x1 images into two classes fake and real one optimizer for each a loss function to train the discriminator create the discriminator create the generator we want to generate 128 coefficients to reshape into a 7x7x128 map here s a featurecomplete gan class overriding compile to use its own signature and implementing the entire gan algorithm in 17 lines in trainstep sample random points in the latent space decode them to fake images combine them with real images assemble labels discriminating real from fake images add random noise to the labels important trick train the discriminator sample random points in the latent space assemble labels that say all real images train the generator note that we should not update the weights of the discriminator update metrics and return their value let s testdrive it prepare the dataset we use both the training test mnist digits create a tensordataset create a dataloader the ideas behind deep learning are simple so why should their implementation be painful introduction when you re doing supervised learning you can use fit and everything works smoothly when you need to take control of every little detail you can write your own training loop entirely from scratch but what if you need a custom training algorithm but you still want to benefit from the convenient features of fit such as callbacks built in distribution support or step fusing a core principle of keras is progressive disclosure of complexity you should always be able to get into lower level workflows in a gradual way you shouldn t fall off a cliff if the high level functionality doesn t exactly match your use case you should be able to gain more control over the small details while retaining a commensurate amount of high level convenience when you need to customize what fit does you should override the training step function of the model class this is the function that is called by fit for every batch of data you will then be able to call fit as usual and it will be running your own learning algorithm note that this pattern does not prevent you from building models with the functional api you can do this whether you re building sequential models functional api models or subclassed models let s see how that works setup this guide can only be run with the torch backend a first simple example let s start from a simple example we create a new class that subclasses keras model we just override the method train_step self data we return a dictionary mapping metric names including the loss to their current value the input argument data is what gets passed to fit as training data if you pass numpy arrays by calling fit x y then data will be the tuple x y if you pass a torch utils data dataloader or a tf data dataset by calling fit dataset then data will be what gets yielded by dataset at each batch in the body of the train_step method we implement a regular training update similar to what you are already familiar with importantly we compute the loss via self compute_loss which wraps the loss es function s that were passed to compile similarly we call metric update_state y y_pred on metrics from self metrics to update the state of the metrics that were passed in compile and we query results from self metrics at the end to retrieve their current value unpack the data its structure depends on your model and on what you pass to fit call torch nn module zero_grad to clear the leftover gradients for the weights from the previous train step compute loss forward pass call torch tensor backward on the loss to compute gradients for the weights update weights update metrics includes the metric that tracks the loss return a dict mapping metric names to current value note that it will include the loss tracked in self metrics let s try this out construct and compile an instance of custommodel just use fit as usual going lower level naturally you could just skip passing a loss function in compile and instead do everything manually in train_step likewise for metrics here s a lower level example that only uses compile to configure the optimizer we start by creating metric instances to track our loss and a mae score in __init__ we implement a custom train_step that updates the state of these metrics by calling update_state on them then query them via result to return their current average value to be displayed by the progress bar and to be pass to any callback note that we would need to call reset_states on our metrics between each epoch otherwise calling result would return an average since the start of training whereas we usually work with per epoch averages thankfully the framework can do that for us just list any metric you want to reset in the metrics property of the model the model will call reset_states on any object listed here at the beginning of each fit epoch or at the beginning of a call to evaluate call torch nn module zero_grad to clear the leftover gradients for the weights from the previous train step compute loss forward pass call torch tensor backward on the loss to compute gradients for the weights update weights compute our own metrics we list our metric objects here so that reset_states can be called automatically at the start of each epoch or at the start of evaluate construct an instance of custommodel we don t passs a loss or metrics here just use fit as usual you can use callbacks etc supporting sample_weight class_weight you may have noticed that our first basic example didn t make any mention of sample weighting if you want to support the fit arguments sample_weight and class_weight you d simply do the following unpack sample_weight from the data argument pass it to compute_loss update_state of course you could also just apply it manually if you don t rely on compile for losses metrics that s it unpack the data its structure depends on your model and on what you pass to fit call torch nn module zero_grad to clear the leftover gradients for the weights from the previous train step compute loss forward pass call torch tensor backward on the loss to compute gradients for the weights update weights update metrics includes the metric that tracks the loss return a dict mapping metric names to current value note that it will include the loss tracked in self metrics construct and compile an instance of custommodel you can now use sample_weight argument providing your own evaluation step what if you want to do the same for calls to model evaluate then you would override test_step in exactly the same way here s what it looks like unpack the data compute predictions updates the metrics tracking the loss update the metrics return a dict mapping metric names to current value note that it will include the loss tracked in self metrics construct an instance of custommodel evaluate with our custom test_step wrapping up an end to end gan example let s walk through an end to end example that leverages everything you just learned let s consider a generator network meant to generate 28x28x1 images a discriminator network meant to classify 28x28x1 images into two classes fake and real one optimizer for each a loss function to train the discriminator create the discriminator create the generator we want to generate 128 coefficients to reshape into a 7x7x128 map here s a feature complete gan class overriding compile to use its own signature and implementing the entire gan algorithm in 17 lines in train_step sample random points in the latent space decode them to fake images combine them with real images assemble labels discriminating real from fake images add random noise to the labels important trick train the discriminator sample random points in the latent space assemble labels that say all real images train the generator note that we should not update the weights of the discriminator update metrics and return their value let s test drive it prepare the dataset we use both the training test mnist digits create a tensordataset create a dataloader the ideas behind deep learning are simple so why should their implementation be painful
import os


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

import torch
import keras
from keras import layers
import numpy as np




class CustomModel(keras.Model):
    def train_step(self, data):
        
        
        x, y = data

        
        
        self.zero_grad()

        
        y_pred = self(x, training=True)  
        loss = self.compute_loss(y=y, y_pred=y_pred)

        
        
        loss.backward()

        trainable_weights = [v for v in self.trainable_weights]
        gradients = [v.value.grad for v in trainable_weights]

        
        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)

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





inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=3)




class CustomModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
        self.loss_fn = keras.losses.MeanSquaredError()

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

        
        
        self.zero_grad()

        
        y_pred = self(x, training=True)  
        loss = self.loss_fn(y, y_pred)

        
        
        loss.backward()

        trainable_weights = [v for v in self.trainable_weights]
        gradients = [v.value.grad for v in trainable_weights]

        
        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        
        self.loss_tracker.update_state(loss)
        self.mae_metric.update_state(y, y_pred)
        return {
            "loss": self.loss_tracker.result(),
            "mae": self.mae_metric.result(),
        }

    @property
    def metrics(self):
        
        
        
        return [self.loss_tracker, self.mae_metric]



inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)


model.compile(optimizer="adam")


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)





class CustomModel(keras.Model):
    def train_step(self, data):
        
        
        if len(data) == 3:
            x, y, sample_weight = data
        else:
            sample_weight = None
            x, y = data

        
        
        self.zero_grad()

        
        y_pred = self(x, training=True)  
        loss = self.compute_loss(
            y=y,
            y_pred=y_pred,
            sample_weight=sample_weight,
        )

        
        
        loss.backward()

        trainable_weights = [v for v in self.trainable_weights]
        gradients = [v.value.grad for v in trainable_weights]

        
        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred, sample_weight=sample_weight)

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



inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
sw = np.random.random((1000, 1))
model.fit(x, y, sample_weight=sw, epochs=3)




class CustomModel(keras.Model):
    def test_step(self, data):
        
        x, y = data
        
        y_pred = self(x, training=False)
        
        loss = self.compute_loss(y=y, y_pred=y_pred)
        
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)
        
        
        return {m.name: m.result() for m in self.metrics}



inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(loss="mse", metrics=["mae"])


x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.evaluate(x, y)




discriminator = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.GlobalMaxPooling2D(),
        layers.Dense(1),
    ],
    name="discriminator",
)


latent_dim = 128
generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        
        layers.Dense(7 * 7 * 128),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Reshape((7, 7, 128)),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
    ],
    name="generator",
)




class GAN(keras.Model):
    def __init__(self, discriminator, generator, latent_dim):
        super().__init__()
        self.discriminator = discriminator
        self.generator = generator
        self.latent_dim = latent_dim
        self.d_loss_tracker = keras.metrics.Mean(name="d_loss")
        self.g_loss_tracker = keras.metrics.Mean(name="g_loss")
        self.seed_generator = keras.random.SeedGenerator(1337)
        self.built = True

    @property
    def metrics(self):
        return [self.d_loss_tracker, self.g_loss_tracker]

    def compile(self, d_optimizer, g_optimizer, loss_fn):
        super().compile()
        self.d_optimizer = d_optimizer
        self.g_optimizer = g_optimizer
        self.loss_fn = loss_fn

    def train_step(self, real_images):
        if isinstance(real_images, tuple):
            real_images = real_images[0]
        
        batch_size = real_images.shape[0]
        random_latent_vectors = keras.random.normal(
            shape=(batch_size, self.latent_dim), seed=self.seed_generator
        )

        
        generated_images = self.generator(random_latent_vectors)

        
        real_images = torch.tensor(real_images)
        combined_images = torch.concat([generated_images, real_images], axis=0)

        
        labels = torch.concat(
            [torch.ones((batch_size, 1)), torch.zeros((batch_size, 1))], axis=0
        )
        
        labels += 0.05 * keras.random.uniform(
            labels.shape, seed=self.seed_generator
        )

        
        self.zero_grad()
        predictions = self.discriminator(combined_images)
        d_loss = self.loss_fn(labels, predictions)
        d_loss.backward()
        grads = [v.value.grad for v in self.discriminator.trainable_weights]
        with torch.no_grad():
            self.d_optimizer.apply(grads, self.discriminator.trainable_weights)

        
        random_latent_vectors = keras.random.normal(
            shape=(batch_size, self.latent_dim), seed=self.seed_generator
        )

        
        misleading_labels = torch.zeros((batch_size, 1))

        
        
        self.zero_grad()
        predictions = self.discriminator(self.generator(random_latent_vectors))
        g_loss = self.loss_fn(misleading_labels, predictions)
        grads = g_loss.backward()
        grads = [v.value.grad for v in self.generator.trainable_weights]
        with torch.no_grad():
            self.g_optimizer.apply(grads, self.generator.trainable_weights)

        
        self.d_loss_tracker.update_state(d_loss)
        self.g_loss_tracker.update_state(g_loss)
        return {
            "d_loss": self.d_loss_tracker.result(),
            "g_loss": self.g_loss_tracker.result(),
        }





batch_size = 64
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
all_digits = np.concatenate([x_train, x_test])
all_digits = all_digits.astype("float32") / 255.0
all_digits = np.reshape(all_digits, (-1, 28, 28, 1))


dataset = torch.utils.data.TensorDataset(
    torch.from_numpy(all_digits), torch.from_numpy(all_digits)
)

dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)

gan = GAN(
    discriminator=discriminator, generator=generator, latent_dim=latent_dim
)
gan.compile(
    d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
    g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
    loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),
)

gan.fit(dataloader, epochs=1)