title writing a training loop from scratch in tensorflow fchollethttps twitter comfchollet date created 20190301 last modified 20230625 description writing lowlevel training evaluation loops in tensorflow accelerator none setup this guide can only be run with the tensorflow backend introduction keras provides default training and evaluation loops fit and evaluate their usage is covered in the guide training evaluation with the builtin methodshttps keras ioguidestrainingwithbuiltinmethods if you want to customize the learning algorithm of your model while still leveraging the convenience of fit for instance to train a gan using fit you can subclass the model class and implement your own trainstep method which is called repeatedly during fit now if you want very lowlevel control over training evaluation you should write your own training evaluation loops from scratch this is what this guide is about a first endtoend example let s consider a simple mnist model let s train it using minibatch gradient with a custom training loop first we re going to need an optimizer a loss function and a dataset instantiate an optimizer instantiate a loss function prepare the training dataset reserve 10 000 samples for validation prepare the training dataset prepare the validation dataset calling a model inside a gradienttape scope enables you to retrieve the gradients of the trainable weights of the layer with respect to a loss value using an optimizer instance you can use these gradients to update these variables which you can retrieve using model trainableweights here s our training loop step by step we open a for loop that iterates over epochs for each epoch we open a for loop that iterates over the dataset in batches for each batch we open a gradienttape scope inside this scope we call the model forward pass and compute the loss outside the scope we retrieve the gradients of the weights of the model with regard to the loss finally we use the optimizer to update the weights of the model based on the gradients iterate over the batches of the dataset open a gradienttape to record the operations run during the forward pass which enables autodifferentiation run the forward pass of the layer the operations that the layer applies to its inputs are going to be recorded on the gradienttape compute the loss value for this minibatch use the gradient tape to automatically retrieve the gradients of the trainable variables with respect to the loss run one step of gradient descent by updating the value of the variables to minimize the loss log every 100 batches lowlevel handling of metrics let s add metrics monitoring to this basic loop you can readily reuse the builtin metrics or custom ones you wrote in such training loops written from scratch here s the flow instantiate the metric at the start of the loop call metric updatestate after each batch call metric result when you need to display the current value of the metric call metric resetstate when you need to clear the state of the metric typically at the end of an epoch let s use this knowledge to compute sparsecategoricalaccuracy on training and validation data at the end of each epoch get a fresh model instantiate an optimizer to train the model instantiate a loss function prepare the metrics here s our training evaluation loop iterate over the batches of the dataset update training metric log every 100 batches display metrics at the end of each epoch reset training metrics at the end of each epoch run a validation loop at the end of each epoch update val metrics speedingup your training step with tf function the default runtime in tensorflow is eager execution as such our training loop above executes eagerly this is great for debugging but graph compilation has a definite performance advantage describing your computation as a static graph enables the framework to apply global performance optimizations this is impossible when the framework is constrained to greedily execute one operation after another with no knowledge of what comes next you can compile into a static graph any function that takes tensors as input just add a tf function decorator on it like this let s do the same with the evaluation step now let s rerun our training loop with this compiled training step iterate over the batches of the dataset log every 100 batches display metrics at the end of each epoch reset training metrics at the end of each epoch run a validation loop at the end of each epoch much faster isn t it lowlevel handling of losses tracked by the model layers models recursively track any losses created during the forward pass by layers that call self addlossvalue the resulting list of scalar loss values are available via the property model losses at the end of the forward pass if you want to be using these loss components you should sum them and add them to the main loss in your training step consider this layer that creates an activity regularization loss let s build a really simple model that uses it insert activity regularization as a layer here s what our training step should look like now add any extra losses created during the forward pass summary now you know everything there is to know about using builtin training loops and writing your own from scratch to conclude here s a simple endtoend example that ties together everything you ve learned in this guide a dcgan trained on mnist digits endtoend example a gan training loop from scratch you may be familiar with generative adversarial networks gans gans can generate new images that look almost real by learning the latent distribution of a training dataset of images the latent space of the images a gan is made of two parts a generator model that maps points in the latent space to points in image space a discriminator model a classifier that can tell the difference between real images from the training dataset and fake images the output of the generator network a gan training loop looks like this 1 train the discriminator sample a batch of random points in the latent space turn the points into fake images via the generator model get a batch of real images and combine them with the generated images train the discriminator model to classify generated vs real images 2 train the generator sample random points in the latent space turn the points into fake images via the generator network get a batch of real images and combine them with the generated images train the generator model to fool the discriminator and classify the fake images as real for a much more detailed overview of how gans works see deep learning with pythonhttps www manning combooksdeeplearningwithpython let s implement this training loop first create the discriminator meant to classify fake vs real digits then let s create a generator network that turns latent vectors into outputs of shape 28 28 1 representing mnist digits we want to generate 128 coefficients to reshape into a 7x7x128 map here s the key bit the training loop as you can see it is quite straightforward the training step function only takes 17 lines instantiate one optimizer for the discriminator and another for the generator instantiate a loss function 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 let s train our gan by repeatedly calling trainstep on batches of images since our discriminator and generator are convnets you re going to want to run this code on a gpu prepare the dataset we use both the training test mnist digits train the discriminator generator on one batch of real images logging print metrics save one generated image to limit execution time we stop after 10 steps remove the lines below to actually train the model that s it you ll get nicelooking fake mnist digits after just 30s of training on the colab gpu setup this guide can only be run with the tensorflow backend introduction keras provides default training and evaluation loops fit and evaluate their usage is covered in the guide training evaluation with the built in methods https keras io guides training_with_built_in_methods if you want to customize the learning algorithm of your model while still leveraging the convenience of fit for instance to train a gan using fit you can subclass the model class and implement your own train_step method which is called repeatedly during fit now if you want very low level control over training evaluation you should write your own training evaluation loops from scratch this is what this guide is about a first end to end example let s consider a simple mnist model let s train it using mini batch gradient with a custom training loop first we re going to need an optimizer a loss function and a dataset instantiate an optimizer instantiate a loss function prepare the training dataset reserve 10 000 samples for validation prepare the training dataset prepare the validation dataset calling a model inside a gradienttape scope enables you to retrieve the gradients of the trainable weights of the layer with respect to a loss value using an optimizer instance you can use these gradients to update these variables which you can retrieve using model trainable_weights here s our training loop step by step we open a for loop that iterates over epochs for each epoch we open a for loop that iterates over the dataset in batches for each batch we open a gradienttape scope inside this scope we call the model forward pass and compute the loss outside the scope we retrieve the gradients of the weights of the model with regard to the loss finally we use the optimizer to update the weights of the model based on the gradients iterate over the batches of the dataset open a gradienttape to record the operations run during the forward pass which enables auto differentiation run the forward pass of the layer the operations that the layer applies to its inputs are going to be recorded on the gradienttape logits for this minibatch compute the loss value for this minibatch use the gradient tape to automatically retrieve the gradients of the trainable variables with respect to the loss run one step of gradient descent by updating the value of the variables to minimize the loss log every 100 batches low level handling of metrics let s add metrics monitoring to this basic loop you can readily reuse the built in metrics or custom ones you wrote in such training loops written from scratch here s the flow instantiate the metric at the start of the loop call metric update_state after each batch call metric result when you need to display the current value of the metric call metric reset_state when you need to clear the state of the metric typically at the end of an epoch let s use this knowledge to compute sparsecategoricalaccuracy on training and validation data at the end of each epoch get a fresh model instantiate an optimizer to train the model instantiate a loss function prepare the metrics here s our training evaluation loop iterate over the batches of the dataset update training metric log every 100 batches display metrics at the end of each epoch reset training metrics at the end of each epoch run a validation loop at the end of each epoch update val metrics speeding up your training step with tf function the default runtime in tensorflow is eager execution as such our training loop above executes eagerly this is great for debugging but graph compilation has a definite performance advantage describing your computation as a static graph enables the framework to apply global performance optimizations this is impossible when the framework is constrained to greedily execute one operation after another with no knowledge of what comes next you can compile into a static graph any function that takes tensors as input just add a tf function decorator on it like this let s do the same with the evaluation step now let s re run our training loop with this compiled training step iterate over the batches of the dataset log every 100 batches display metrics at the end of each epoch reset training metrics at the end of each epoch run a validation loop at the end of each epoch much faster isn t it low level handling of losses tracked by the model layers models recursively track any losses created during the forward pass by layers that call self add_loss value the resulting list of scalar loss values are available via the property model losses at the end of the forward pass if you want to be using these loss components you should sum them and add them to the main loss in your training step consider this layer that creates an activity regularization loss let s build a really simple model that uses it insert activity regularization as a layer here s what our training step should look like now add any extra losses created during the forward pass summary now you know everything there is to know about using built in training loops and writing your own from scratch to conclude here s a simple end to end example that ties together everything you ve learned in this guide a dcgan trained on mnist digits end to end example a gan training loop from scratch you may be familiar with generative adversarial networks gans gans can generate new images that look almost real by learning the latent distribution of a training dataset of images the latent space of the images a gan is made of two parts a generator model that maps points in the latent space to points in image space a discriminator model a classifier that can tell the difference between real images from the training dataset and fake images the output of the generator network a gan training loop looks like this 1 train the discriminator sample a batch of random points in the latent space turn the points into fake images via the generator model get a batch of real images and combine them with the generated images train the discriminator model to classify generated vs real images 2 train the generator sample random points in the latent space turn the points into fake images via the generator network get a batch of real images and combine them with the generated images train the generator model to fool the discriminator and classify the fake images as real for a much more detailed overview of how gans works see deep learning with python https www manning com books deep learning with python let s implement this training loop first create the discriminator meant to classify fake vs real digits then let s create a generator network that turns latent vectors into outputs of shape 28 28 1 representing mnist digits we want to generate 128 coefficients to reshape into a 7x7x128 map here s the key bit the training loop as you can see it is quite straightforward the training step function only takes 17 lines instantiate one optimizer for the discriminator and another for the generator instantiate a loss function 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 let s train our gan by repeatedly calling train_step on batches of images since our discriminator and generator are convnets you re going to want to run this code on a gpu prepare the dataset we use both the training test mnist digits in practice you need at least 20 epochs to generate nice digits train the discriminator generator on one batch of real images logging print metrics save one generated image to limit execution time we stop after 10 steps remove the lines below to actually train the model that s it you ll get nice looking fake mnist digits after just 30s of training on the colab gpu
import time
import os


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

import tensorflow as tf
import keras
import numpy as np






def get_model():
    inputs = keras.Input(shape=(784,), name="digits")
    x1 = keras.layers.Dense(64, activation="relu")(inputs)
    x2 = keras.layers.Dense(64, activation="relu")(x1)
    outputs = keras.layers.Dense(10, name="predictions")(x2)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


model = get_model()




optimizer = keras.optimizers.Adam(learning_rate=1e-3)

loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)


batch_size = 32
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = np.reshape(x_train, (-1, 784))
x_test = np.reshape(x_test, (-1, 784))


x_val = x_train[-10000:]
y_val = y_train[-10000:]
x_train = x_train[:-10000]
y_train = y_train[:-10000]


train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)


val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(batch_size)



epochs = 3
for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")

    
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        
        
        with tf.GradientTape() as tape:
            
            
            
            
            logits = model(
                x_batch_train, training=True
            )  

            
            loss_value = loss_fn(y_batch_train, logits)

        
        
        grads = tape.gradient(loss_value, model.trainable_weights)

        
        
        optimizer.apply(grads, model.trainable_weights)

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {float(loss_value):.4f}"
            )
            print(f"Seen so far: {(step + 1) * batch_size} samples")




model = get_model()


optimizer = keras.optimizers.Adam(learning_rate=1e-3)

loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)


train_acc_metric = keras.metrics.SparseCategoricalAccuracy()
val_acc_metric = keras.metrics.SparseCategoricalAccuracy()



epochs = 2
for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    start_time = time.time()

    
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            logits = model(x_batch_train, training=True)
            loss_value = loss_fn(y_batch_train, logits)
        grads = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply(grads, model.trainable_weights)

        
        train_acc_metric.update_state(y_batch_train, logits)

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {float(loss_value):.4f}"
            )
            print(f"Seen so far: {(step + 1) * batch_size} samples")

    
    train_acc = train_acc_metric.result()
    print(f"Training acc over epoch: {float(train_acc):.4f}")

    
    train_acc_metric.reset_state()

    
    for x_batch_val, y_batch_val in val_dataset:
        val_logits = model(x_batch_val, training=False)
        
        val_acc_metric.update_state(y_batch_val, val_logits)
    val_acc = val_acc_metric.result()
    val_acc_metric.reset_state()
    print(f"Validation acc: {float(val_acc):.4f}")
    print(f"Time taken: {time.time() - start_time:.2f}s")




@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply(grads, model.trainable_weights)
    train_acc_metric.update_state(y, logits)
    return loss_value





@tf.function
def test_step(x, y):
    val_logits = model(x, training=False)
    val_acc_metric.update_state(y, val_logits)




epochs = 2
for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    start_time = time.time()

    
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        loss_value = train_step(x_batch_train, y_batch_train)

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {float(loss_value):.4f}"
            )
            print(f"Seen so far: {(step + 1) * batch_size} samples")

    
    train_acc = train_acc_metric.result()
    print(f"Training acc over epoch: {float(train_acc):.4f}")

    
    train_acc_metric.reset_state()

    
    for x_batch_val, y_batch_val in val_dataset:
        test_step(x_batch_val, y_batch_val)

    val_acc = val_acc_metric.result()
    val_acc_metric.reset_state()
    print(f"Validation acc: {float(val_acc):.4f}")
    print(f"Time taken: {time.time() - start_time:.2f}s")






class ActivityRegularizationLayer(keras.layers.Layer):
    def call(self, inputs):
        self.add_loss(1e-2 * tf.reduce_sum(inputs))
        return inputs




inputs = keras.Input(shape=(784,), name="digits")
x = keras.layers.Dense(64, activation="relu")(inputs)

x = ActivityRegularizationLayer()(x)
x = keras.layers.Dense(64, activation="relu")(x)
outputs = keras.layers.Dense(10, name="predictions")(x)

model = keras.Model(inputs=inputs, outputs=outputs)




@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss_value = loss_fn(y, logits)
        
        loss_value += sum(model.losses)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply(grads, model.trainable_weights)
    train_acc_metric.update_state(y, logits)
    return loss_value






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



latent_dim = 128

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




d_optimizer = keras.optimizers.Adam(learning_rate=0.0003)
g_optimizer = keras.optimizers.Adam(learning_rate=0.0004)


loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)


@tf.function
def train_step(real_images):
    
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    
    generated_images = generator(random_latent_vectors)
    
    combined_images = tf.concat([generated_images, real_images], axis=0)

    
    labels = tf.concat(
        [tf.ones((batch_size, 1)), tf.zeros((real_images.shape[0], 1))], axis=0
    )
    
    labels += 0.05 * tf.random.uniform(labels.shape)

    
    with tf.GradientTape() as tape:
        predictions = discriminator(combined_images)
        d_loss = loss_fn(labels, predictions)
    grads = tape.gradient(d_loss, discriminator.trainable_weights)
    d_optimizer.apply(grads, discriminator.trainable_weights)

    
    random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
    
    misleading_labels = tf.zeros((batch_size, 1))

    
    
    with tf.GradientTape() as tape:
        predictions = discriminator(generator(random_latent_vectors))
        g_loss = loss_fn(misleading_labels, predictions)
    grads = tape.gradient(g_loss, generator.trainable_weights)
    g_optimizer.apply(grads, generator.trainable_weights)
    return d_loss, g_loss, generated_images





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 = tf.data.Dataset.from_tensor_slices(all_digits)
dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)

epochs = 1  
save_dir = "./"

for epoch in range(epochs):
    print(f"\nStart epoch {epoch}")

    for step, real_images in enumerate(dataset):
        
        d_loss, g_loss, generated_images = train_step(real_images)

        
        if step % 100 == 0:
            
            print(f"discriminator loss at step {step}: {d_loss:.2f}")
            print(f"adversarial loss at step {step}: {g_loss:.2f}")

            
            img = keras.utils.array_to_img(
                generated_images[0] * 255.0, scale=False
            )
            img.save(os.path.join(save_dir, f"generated_img_{step}.png"))

        
        
        if step > 10:
            break