Deep Learning with PyTorch Step-by-Step: Part I - Fundamentals/

...

/

Rethinking the Training Loop

Rethinking the Training Loop

Learn how you can reduce the boilerplate section from the training loop by using higher-order functions.

We'll cover the following...

Training step

As already mentioned, the higher-order function that builds a training step function for us is taking the key elements of our training loop: model, loss, and optimizer. The actual training step function to be returned will have two arguments, namely, features and labels, and will return the corresponding loss value.

Creating the higher-order function for training step

Apart from returning the loss value, the inner perform_train_step() function below is the same as the code inside the loop in model training V0. The code should look like this:

Press + to interact
def make_train_step(model, loss_fn, optimizer):
    # Builds function that performs a step in the train loop
    def perform_train_step(x, y):
        # Sets model to TRAIN mode
        model.train()
        
        # Step 1 - computes model's predictions - forward pass
        yhat = model(x)
        # Step 2 - computes the loss
        loss = loss_fn(yhat, y)
        # Step 3 - computes gradients for "b" and "w" parameters
        loss.backward()
        # Step 4 - updates parameters using gradients and
        # the learning rate
        optimizer.step()
        optimizer.zero_grad()
        
        # Returns the loss
        return loss.item()
    
    # Returns the function that will be called inside the 
    # train loop
    return perform_train_step

Updating model configuration code

Then, we need to update our model configuration code to call this higher-order function to build a train_step function. But we need to run the data preparation script ...

Access this course and 1400+ top-rated courses and projects.