Wondering why my model could be plateauing

aadithyaa: <p>I’m training a deep neural network on temporal graph data. Currently, I’m trying to get a feel for how large / complex of a model I should aim for, so I’m trying to overfit to my smallest dataset. Here is my code.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
import sys
import random
import os
import datetime

# Command to run
# python3 my_dnn.py [dataset_basename] [num_layers] [batch_size] [neurons_per_hidden_layer] [data_directory] [num_epochs] [learning_rate] [rng_seed]
# python3 my_dnn.py enron 25 20 256 ./formatted_data 10 0.01 48444316

# Setting the CUBLAS environment variable for deterministic operations
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'

# Name of the dataset to train on
DATASET_BASENAME = sys.argv[1]

# Number of features (two endpoints of an edge + insertion time and two numeric features)
INPUT_SIZE = 5  

# Number of stacked rnn layers
NUM_LAYERS = int(sys.argv[2])
BATCH_SIZE = int(sys.argv[3])
OUTPUT_SIZE = 1
HIDDEN_DIM = int(sys.argv[4])
DATA_DIRECTORY = str(sys.argv[5])

# Number of previous timestamps taken into account
SEQ_LENGTH = 2
NUM_EPOCHS = int(sys.argv[6])
LEARNING_RATE = float(sys.argv[7])
LSTM = False

# Try to use gpu if it's there
is_cuda = torch.cuda.is_available()
device = torch.device("cuda" if is_cuda else "cpu")
print("GPU is available" if is_cuda else "GPU not available, CPU used")


# We require all algorithms to be deterministic and fix random seed 48444316 for reproducibility
torch.use_deterministic_algorithms(True)
def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
set_seed(int(sys.argv[8]))


class TemporalGraphDataset(Dataset):
    def __init__(self, feature_file, label_file, seq_length):
        self.features = pd.read_csv(feature_file)
        self.labels = pd.read_csv(label_file)
        self.seq_length = seq_length

    def __len__(self):
        return len(self.features) - self.seq_length + 1

    def __getitem__(self, idx):
        start_idx = idx
        end_idx = idx + self.seq_length
        data = self.features.iloc[start_idx:end_idx].to_numpy()
        target = self.labels.iloc[end_idx - 1].to_numpy()
        return torch.tensor(data, dtype=torch.float32), torch.tensor(target, dtype=torch.float32)

features_path = f'{DATA_DIRECTORY}/formatted_{DATASET_BASENAME}/{DATASET_BASENAME}_features.csv'
labels_path = f'{DATA_DIRECTORY}/formatted_{DATASET_BASENAME}/{DATASET_BASENAME}_labels.csv'
dataset = TemporalGraphDataset(features_path, labels_path, SEQ_LENGTH)

train_dataset, val_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 100, 100])

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=1)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=1)

class DNN(nn.Module):
    def __init__(self, input_size, output_size, hidden_dim, n_layers):
        super(DNN, self).__init__()
        self.rnn = nn.RNN(input_size, hidden_dim, n_layers, batch_first=True) if not LSTM else nn.LSTM(input_size, hidden_dim, n_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_size)

    def forward(self, x):
        hidden = self.init_hidden(x.size(0))
        out, _ = self.rnn(x) if not LSTM else self.rnn(x, hidden)
        return self.fc(out[:, -1, :])  # Taking the last timestep output

    def init_hidden(self, batch_size):
        return torch.zeros(NUM_LAYERS, batch_size, HIDDEN_DIM, device=device)

model = DNN(INPUT_SIZE, OUTPUT_SIZE, HIDDEN_DIM, NUM_LAYERS).to(device)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.MSELoss()

# Validation function
def validate(model, dataloader, criterion):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for data, target in dataloader:
            model.zero_grad()
            data, target = data.cuda().float(), target.cuda().float()
            output = model(data)
            loss = criterion(output, target)
            total_loss += loss.item()
    return total_loss / len(dataloader)

# Training function with validation after each
def train_model(model, train_loader, val_loader, optimizer, criterion, epochs):
    train_losses = []
    val_losses = []
    for epoch in range(epochs):
        model.train()
        total_train_loss = 0
        for data, target in train_loader:
            model.zero_grad()
            data, target = data.cuda(), target.cuda()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm = 1000)
            optimizer.step()
            total_train_loss += loss.item()
            
        # Store average train loss over all examples
        train_losses.append(total_train_loss / len(train_loader))

        val_loss = validate(model, val_loader, criterion)
        val_losses.append(val_loss)

        print(f'Epoch {epoch+1}, Training Loss: {train_losses[-1]}, Validation Loss: {val_losses[-1]}')

    # Plotting
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label='Training Loss')
    plt.plot(val_losses, label='Validation Loss')
    plt.title('Training and Validation Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    today = datetime.datetime.now().strftime("%Y-%m-%d")
    plot_filename = f"{DATASET_BASENAME}_{today}.png"
    plt.savefig(plot_filename)
    print(f"Plot saved as {plot_filename}")
    plt.show()

train_model(model, train_loader, val_loader, optimizer, criterion, NUM_EPOCHS)
</code></pre>
<p>Here is the output from the terminal</p>
<pre><code class="lang-auto">(env) python3 my_dnn.py enron 25 20 256 ./formatted_data 10 0.01 48444316
GPU is available
Epoch 1, Training Loss: 255494918.86796424, Validation Loss: 178531414.4
Epoch 2, Training Loss: 130291595.62125197, Validation Loss: 80107624.8
Epoch 3, Training Loss: 52912246.85954761, Validation Loss: 29422431.2
Epoch 4, Training Loss: 22821705.16570226, Validation Loss: 21934850.0
Epoch 5, Training Loss: 20782540.18095739, Validation Loss: 21980772.4
Epoch 6, Training Loss: 20782165.60441873, Validation Loss: 22031612.8
Epoch 7, Training Loss: 20781856.248290375, Validation Loss: 21950381.0
Epoch 8, Training Loss: 20781392.246712256, Validation Loss: 21956635.4
Epoch 9, Training Loss: 20776616.859021567, Validation Loss: 21970061.2
Epoch 10, Training Loss: 20781443.142556548, Validation Loss: 21955459.8
</code></pre>
<p>It seems to plateau at ~2 * 10^7 mean squared error. Why could this be? Are there any obvious bugs in my code? What are some ideas on how I can get out of this plateau?</p>