Post cover image
https://blogs.sap.com/2020/05/07/advanced-debugging-of-mta-based-apps/

September 7, 2023

How to Debug a Machine Learning Model: My Experience with LSTMs as an Example

In machine learning, building a model is just one part of the puzzle. Some basic models are pre-constructed, but the real challenge often…

By Renee LIN

6 min read

In machine learning, building a model is just one part of the puzzle. Some basic models are pre-constructed, but the real challenge often arises when a model doesn't perform as expected. There might be issues with accuracy, or perhaps the model doesn't handle unseen data well. Debugging then becomes indispensable. While I'm not an AI expert, the process of building and training models was quite daunting for me. I might not have the expertise to offer, but I can share experiences that many beginners might encounter. I'd like to share what I've learned from the challenges of building a simple LSTM model. The main checkpoints are:

  1. The implementation
  2. Inputs and outputs during the process
  3. The loss during the training
  4. Regularization techniques
  5. Hyperparameters tuning

0. A little background of my project

I'm working on a car-following model that takes vehicle speed, speed difference, and spacing to the leading vehicle as inputs and predicts acceleration. The model is evaluated by the root mean square error (RMSE) of spacing in the trajectories from the dataset. I have successfully applied imitation models, specifically behaviour cloning and generative adversarial imitation learning, to the dataset and achieved an accuracy of around 6 meters. This means that, on average, the generated location of the car is similar to those in the dataset with a deviation of approximately 6 meters. However, my problem is that my LSTM model only achieved an accuracy of 16 meters, which is definitely not correct. According to a benchmark paper[1], the expected accuracy is around 5 meters (the root of MSE being 23.86).

None

What is wrong?

1. The implementation

The LSTM model has been implemented in several popular frameworks, such as PyTorch. I'm using a simple GRU model at the bottom of this post. Upon examining my code, I discovered an obvious mistake: I hadn't trained the model using mini-batches. This oversight caused my computer to slow down during training. The Activity Monitor indicated that Python was consuming around 10GB of memory. When I transferred the process to Colab and tried to address the problem by utilizing a GPU, I received an 'out of memory' error. However, after setting a batch size of 32, the error was resolved, and Python used only 1.5GB of memory on my Mac.

OutOfMemoryError: CUDA out of memory. Tried to allocate 15.98 GiB (GPU 0; 14.75 GiB total capacity; 1.28 GiB already allocated; 12.65 GiB free; 1.29 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Another mistake I made was assuming that functions like "evaluation_all" were correct, as I had used them in other models and had already debugged them. However, with LSTM, we make predictions based on some historical steps. Initially, we need to reserve that sequence length, but I forgot. I printed out the generated trajectory and compared it to the ones from the dataset and noticed that the first prediction leapt. This printout and examination step is what I want to emphasize next.

ini_sub_pos = expert_traj[i]['subject_loc'].loc[0] #should be .loc[seq_len]

2. Check the input and output during the process

The input should be 20 or 30 steps of vehicle speed, leading vehicle speed, speed difference and spacing. I've scaled these input so that they all fall between -1 and 1. Due to this scaling, it's also necessary to check the inverse values at certain points. Based on my experience, the model should output an acceleration value between -3 and 3."

None

3. Check the loss during the process

Recording and visualizing the loss value during training is standard practice. I often use 'weights and bias,' which facilitates easy logging and parameter sweeping(How to Tune Hyperparameters with Weights and Biases). Since there's only one component here with a single MSE loss, I didn't use it for quick debugging. The model's loss is logged every 10 epochs. I observed that it's decreasing and also noticed the typical 'elbow' pattern, indicating it likely converged.

None

4. Regularization techniques

Apparently, my model is overfitting. The loss is extremely small but the accuracy is 147 meters. The model should be able to generalize to new data, not performing well in a small set but poorly on the testing set.

  • Increased the batch size from 32 to 2048, led to 43 avg_rmse_space
  • I added dropout=0.1, the accuracy became 26

Dropout: During training, we omit a percentage of the neurons from a specific layer. This prevents the neurons from becoming too specialized. There are many other regularization techniques. Since I'm not well-versed in this area, I'd like to learn more and summarize my findings in another post in the near future.

None None

5. Hyperparameters tuning

It's worth mentioning that I upgraded from a free Colab account to a monthly subscription because the free version only provides limited usage of the GPU(not stable, sometimes no access). After opting for the better GPU and higher RAM at 15 AUD/month, the training time dropped from about 40 minutes to 10 minutes. This reduction in time allowed me to experiment with several sets of hyperparameters faster.

Below is my starting point, based on the previous exploration steps.

# Constants
SQUENCE_LEN = 20
INPUT_SIZE = 4

# Hyperparameters and Configuration
NUM_EPOCHS = 200
LEARNING_RATE = 0.0001
HIDDEN_SIZE = 64
NUM_LAYERS = 2
BATCH_SIZE = 2048

After conducting a manual grid search, I settled on this set, achieving an accuracy (RMSE of spacing) of 6 meters. In fact, my primary issue was the implementation mistakes I mentioned earlier. By trying some common values, such as 0.0003, I already achieved good accuracy. It's worth noting that using a random seed and recording the randomization choices assists in grid search as it allows for model reproducibility.

Summary

It's valuable to practice constructing and fine-tuning basic models or utilizing pre-built ones. In real-world scenarios, there's often no need to build complex models ourselves. Make sure the implementation is correct and print the values and loss during the process. While fine-tuning the model, consider some regularization techniques and use tools to sweep the hyperparameters.

[1] Chen, X., Zhu, M., Chen, K., Wang, P., Lu, H., Zhong, H., … & Wang, Y. (2023). FollowNet: A Comprehensive Benchmark for Car-Following Behavior Modeling. arXiv preprint arXiv:2306.05381.

My codes before correcting:

import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from sklearn.model_selection import train_test_split
from data import expert_traj_diff_train, expert_traj_diff_test,expert_traj_test
from lstm_evaluate import evaluate_all
from datetime import date
import os
import time
import matplotlib.pyplot as plt

# Constants
SQUENCE_LEN = 10
INPUT_SIZE = 4

# Hyperparameters and Configuration
NUM_EPOCHS = 1000
LEARNING_RATE = 0.0001
HIDDEN_SIZE = 128
NUM_LAYERS = 1

def sliding_windows(data, seq_length):
    x, y = [], []
    for traj in data:
        for i in range(len(traj)-seq_length-1):
            _x = traj.loc[i:(i+seq_length)][['subject_speed','preceding_speed','v_diff','pos_diff']].values
            _y = traj.loc[i+seq_length][['v_Acc']].values
            x.append(_x)
            y.append(_y)
    return np.array(x), np.array(y)

x, y = sliding_windows(expert_traj_diff_train, SQUENCE_LEN)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2)

trainX = Variable(torch.Tensor(x_train))
trainY = Variable(torch.Tensor(y_train))
testX = Variable(torch.Tensor(x_test))
testY = Variable(torch.Tensor(y_test))

class GRUModel(nn.Module):
    def __init__(self, output_dim, input_dim, hidden_dim, num_layers):
        super(GRUModel, self).__init__()

        # Defining the number of layers and the nodes in each layer
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim

        # GRU layers
        self.gru = nn.GRU(
            input_dim, hidden_dim, num_layers, batch_first=True,
        )

        # Fully connected layer
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # Initializing hidden state for first input with zeros
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()

        # Forward propagation by passing in the input and hidden state into the model
        out, _ = self.gru(x, h0.detach())

        # Reshaping the outputs in the shape of (batch_size, seq_length, hidden_size)
        # so that it can fit into the fully connected layer
        out = out[:, -1, :]

        # Convert the final state to our desired output shape (batch_size, output_dim)
        out = self.fc(out)

        return out

def train(trainX, trainY):
    model_dir = f"models/LSTM-{date.today()}-{time.strftime('%H-%M-%S', time.localtime())}-{LEARNING_RATE}"
    os.makedirs(model_dir, exist_ok=True)

    lstm = GRUModel(1, INPUT_SIZE, HIDDEN_SIZE, NUM_LAYERS)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    lstm = lstm.to(device)
    trainX, trainY = trainX.to(device), trainY.to(device)

    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(lstm.parameters(), lr=LEARNING_RATE)

    losses = []

    for epoch in range(NUM_EPOCHS):
        outputs = lstm(trainX)
        optimizer.zero_grad()

        loss = criterion(outputs, trainY)
        losses.append(loss.item())

        loss.backward()
        optimizer.step()

        if epoch % 10 == 0:
            print(f"Epoch: {epoch}, loss: {loss.item():.5f}")
            torch.save(lstm.state_dict(), f"{model_dir}/{epoch}")

    # Plot the losses over epochs
    plt.plot(losses)
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.title("Training Loss over Epochs")
    plt.show()

    # Evaluate the model
    avg_rmse_speed, avg_rmse_space, negative_v, crash = evaluate_all(expert_traj_diff_test, expert_traj_test, lstm)
    print('avg_rmse_speed',avg_rmse_speed)
    print('avg_rmse_space',avg_rmse_space)
    print('negative_v',negative_v)
    print('crash',crash)

train(trainX, trainY)