title writing a training loop from scratch in pytorch fchollethttps twitter comfchollet date created 20230625 last modified 20230625 description writing lowlevel training evaluation loops in pytorch accelerator none setup this guide can only be run with the torch 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 to write a custom training loop we need the following ingredients a model to train of course an optimizer you could either use a keras optimizers optimizer or a native pytorch optimizer from torch optim a loss function you could either use a keras losses loss or a native pytorch loss from torch nn a dataset you could use any format a tf data dataset a pytorch dataloader a python generator etc let s line them up we ll use torchnative objects in each case except of course for the keras model first let s get the model and the mnist dataset let s consider a simple mnist model create load up the mnist dataset and put it in a torch dataloader prepare the training dataset reserve 10 000 samples for validation create torch datasets create dataloaders for the datasets next here s our pytorch optimizer and our pytorch loss function instantiate a torch optimizer instantiate a torch loss function let s train our model using minibatch gradient with a custom training loop calling loss backward on a loss tensor triggers backpropagation once that s done your optimizer is magically aware of the gradients for each variable and can update its variables which is done via optimizer step tensors variables optimizers are all interconnected to one another via hidden global state also don t forget to call model zerograd before loss backward or you won t get the right gradients for your variables 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 call the model on the input data to retrive the predictions then we use them to compute a loss value we call loss backward to 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 forward pass backward pass optimizer variable updates log every 100 batches as an alternative let s look at what the loop looks like when using a keras optimizer and a keras loss function important differences you retrieve the gradients for the variables via v value grad called on each trainable variable you update your variables via optimizer apply which must be called in a torch nograd scope also a big gotcha while all numpytensorflowjaxkeras apis as well as python unittest apis use the argument order convention fnytrue ypred reference values first predicted values second pytorch actually uses fnypred ytrue for its losses so make sure to invert the order of logits and targets forward pass backward pass call torch tensor backward on the loss to compute gradients for the weights update weights log every 100 batches lowlevel handling of metrics let s add metrics monitoring to this basic training loop you can readily reuse builtin keras 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 categoricalaccuracy 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 forward pass backward pass call torch tensor backward on the loss to compute gradients for the weights update weights 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 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 loop should look like now get a fresh model instantiate an optimizer to train the model instantiate a loss function prepare the metrics forward pass backward pass call torch tensor backward on the loss to compute gradients for the weights update weights 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 that s it setup this guide can only be run with the torch 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 to write a custom training loop we need the following ingredients a model to train of course an optimizer you could either use a keras optimizers optimizer or a native pytorch optimizer from torch optim a loss function you could either use a keras losses loss or a native pytorch loss from torch nn a dataset you could use any format a tf data dataset a pytorch dataloader a python generator etc let s line them up we ll use torch native objects in each case except of course for the keras model first let s get the model and the mnist dataset let s consider a simple mnist model create load up the mnist dataset and put it in a torch dataloader prepare the training dataset reserve 10 000 samples for validation create torch datasets create dataloaders for the datasets next here s our pytorch optimizer and our pytorch loss function instantiate a torch optimizer instantiate a torch loss function let s train our model using mini batch gradient with a custom training loop calling loss backward on a loss tensor triggers backpropagation once that s done your optimizer is magically aware of the gradients for each variable and can update its variables which is done via optimizer step tensors variables optimizers are all interconnected to one another via hidden global state also don t forget to call model zero_grad before loss backward or you won t get the right gradients for your variables 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 call the model on the input data to retrive the predictions then we use them to compute a loss value we call loss backward to 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 forward pass backward pass optimizer variable updates log every 100 batches as an alternative let s look at what the loop looks like when using a keras optimizer and a keras loss function important differences you retrieve the gradients for the variables via v value grad called on each trainable variable you update your variables via optimizer apply which must be called in a torch no_grad scope also a big gotcha while all numpy tensorflow jax keras apis as well as python unittest apis use the argument order convention fn y_true y_pred reference values first predicted values second pytorch actually uses fn y_pred y_true for its losses so make sure to invert the order of logits and targets forward pass backward pass call torch tensor backward on the loss to compute gradients for the weights update weights log every 100 batches low level handling of metrics let s add metrics monitoring to this basic training loop you can readily reuse built in keras 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 categoricalaccuracy 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 forward pass backward pass call torch tensor backward on the loss to compute gradients for the weights update weights 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 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 loop should look like now get a fresh model instantiate an optimizer to train the model instantiate a loss function prepare the metrics forward pass backward pass call torch tensor backward on the loss to compute gradients for the weights update weights 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 that s it
import os


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

import torch
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




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


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


train_dataset = torch.utils.data.TensorDataset(
    torch.from_numpy(x_train), torch.from_numpy(y_train)
)
val_dataset = torch.utils.data.TensorDataset(
    torch.from_numpy(x_val), torch.from_numpy(y_val)
)


train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True
)
val_dataloader = torch.utils.data.DataLoader(
    val_dataset, batch_size=batch_size, shuffle=False
)




model = get_model()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)


loss_fn = torch.nn.CrossEntropyLoss()



epochs = 3
for epoch in range(epochs):
    for step, (inputs, targets) in enumerate(train_dataloader):
        
        logits = model(inputs)
        loss = loss_fn(logits, targets)

        
        model.zero_grad()
        loss.backward()

        
        optimizer.step()

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {loss.detach().numpy():.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.CategoricalCrossentropy(from_logits=True)

for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    for step, (inputs, targets) in enumerate(train_dataloader):
        
        logits = model(inputs)
        loss = loss_fn(targets, logits)

        
        model.zero_grad()
        trainable_weights = [v for v in model.trainable_weights]

        
        
        loss.backward()
        gradients = [v.value.grad for v in trainable_weights]

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

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {loss.detach().numpy():.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.CategoricalCrossentropy(from_logits=True)


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



for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    for step, (inputs, targets) in enumerate(train_dataloader):
        
        logits = model(inputs)
        loss = loss_fn(targets, logits)

        
        model.zero_grad()
        trainable_weights = [v for v in model.trainable_weights]

        
        
        loss.backward()
        gradients = [v.value.grad for v in trainable_weights]

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

        
        train_acc_metric.update_state(targets, logits)

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {loss.detach().numpy():.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_dataloader:
        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}")





class ActivityRegularizationLayer(keras.layers.Layer):
    def call(self, inputs):
        self.add_loss(1e-2 * torch.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)




model = get_model()


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

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


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

for epoch in range(epochs):
    print(f"\nStart of epoch {epoch}")
    for step, (inputs, targets) in enumerate(train_dataloader):
        
        logits = model(inputs)
        loss = loss_fn(targets, logits)
        if model.losses:
            loss = loss + torch.sum(*model.losses)

        
        model.zero_grad()
        trainable_weights = [v for v in model.trainable_weights]

        
        
        loss.backward()
        gradients = [v.value.grad for v in trainable_weights]

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

        
        train_acc_metric.update_state(targets, logits)

        
        if step % 100 == 0:
            print(
                f"Training loss (for 1 batch) at step {step}: {loss.detach().numpy():.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_dataloader:
        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}")