Overfitting issue with video dataset while using PyG

[BhuyanMonami]

I am using Graph CNNs for action recognition for a video dataset. Specifically, as a starting point I have modelled my data in the following way:

1. I have used an off the shelf object detection model to detect objects (with their bounding box coordinates and CNN features) from each frame.
2. I have constructed graph objects with 30 frames and a overlapping window of 15 frames. Thus, each graph object is made up of the different detected objects which are the nodes. As for the edges, I have calculated the IOU between the bounding box coordinates and connected the nodes.
3. In this regard, I have two different types of graphs: spatial(where I connected all objects to each other) and temporal (where I connected only objects of the same type across the 30 frames)
4. Finally I used the HeteroConv wrapper along with GCNConv for my model.

Now, while my model initially seemingly performed "well" on the training and validation data (with an accuracy >90% for each action class), it wasn't able to generalise on unseen testing data. This let me to conclude to my model is overfitting. So as per #2008 I tried using dropout layers after each conv. layer, however, after having done so, the validation accuracy doesn't improve beyond 58%(ish) after using early stopping.

I would be very grateful if some input and suggestions were thrown my way. For example, should I use a different convolution layer? Or different hyperparameters? Or different features etc?

Here is my model:

import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
import math
import pdb
from torch_geometric.nn import Linear,GraphConv,HeteroConv,GCNConv
from torch.nn import BatchNorm1d
from torch_geometric.typing import torch_scatter
from torch_geometric.nn import global_mean_pool
from torch_geometric.nn import global_add_pool
from torch_geometric.nn import global_max_pool

class RGCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
        super(RGCN, self).__init__()
        self.convs1 = torch.nn.ModuleList()
        self.convs2 = torch.nn.ModuleList()
        self.convs3 = torch.nn.ModuleList()
        
        self.batch_norms1 = torch.nn.ModuleList()
        self.batch_norms2 = torch.nn.ModuleList()
        self.batch_norms3 = torch.nn.ModuleList()
        
        self.lin1 = Linear(hidden_channels*3, hidden_channels*3)
        self.lin2 = Linear(hidden_channels*3, out_channels)
        
        self.dropout = nn.Dropout(p=0.5)
        
        for _ in range(num_layers):
            conv1 = HeteroConv({('frame_window', 'spatial', 'frame_window'): GCNConv(in_channels, hidden_channels),
                                ('frame_window', 'temporal', 'frame_window'): GCNConv(in_channels, hidden_channels)}, aggr='sum')
            conv2 = HeteroConv({('frame_window', 'spatial', 'frame_window'): GCNConv(hidden_channels, hidden_channels),
                                ('frame_window', 'temporal', 'frame_window'): GCNConv(hidden_channels, hidden_channels)}, aggr='sum')
            conv3 = HeteroConv({('frame_window', 'spatial', 'frame_window'): GCNConv(hidden_channels, hidden_channels),
                                ('frame_window', 'temporal', 'frame_window'): GCNConv(hidden_channels, hidden_channels)}, aggr='sum')
            bn1 = nn.BatchNorm1d(hidden_channels)
            bn2 = nn.BatchNorm1d(hidden_channels)
            bn3 = nn.BatchNorm1d(hidden_channels)
            self.convs1.append(conv1)
            self.convs2.append(conv2)
            self.convs3.append(conv3)
            self.batch_norms1.append(bn1)
            self.batch_norms2.append(bn2)
            self.batch_norms3.append(bn3)      


    def forward(self, data):
        x_dict = data.x_dict
        edge_index_dict = data.edge_index_dict
        edge_weight_dict = data.edge_weight_dict
        batch = data['frame_window'].batch
        
        for conv1, conv2, conv3, bn1, bn2, bn3 in zip(self.convs1, self.convs2, self.convs3, 
                                                      self.batch_norms1, self.batch_norms2, self.batch_norms3):
            gc1 = conv1(x_dict, edge_index_dict, edge_weight_dict)
            gc1 = {key: bn1(x.relu()) for key, x in gc1.items()}
            gc2 = conv2(gc1, edge_index_dict, edge_weight_dict)
            gc2 = {key: self.dropout(bn2(x.relu())) for key, x in gc2.items()}
            gc3 = conv3(gc2, edge_index_dict, edge_weight_dict)
            gc3 = {key: self.dropout(bn3(x.relu())) for key, x in gc3.items()}

        out1 = global_mean_pool(gc3['frame_window'], batch)
        out2 = global_add_pool(gc3['frame_window'], batch)
        out3 = global_max_pool(gc3['frame_window'], batch) 
        out = torch.cat((out1, out2, out3), dim=1)
        out = F.relu(self.lin1(out))
        out = F.dropout(out, p=0.5, training=self.training)
        out = self.lin2(out)
        return out

And here is my trainer:

import datetime
import torch_scatter
from sklearn.metrics import accuracy_score
from sklearn import metrics
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

lr = 0.00001
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.7, patience=2,
                                                              min_lr=1e-7)
class_weights = class_weights.to(device)
loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights)
all_preds = []
all_labels = []

def train(epochs):
    early_stopper = EarlyStopper()
    global all_preds, all_labels
    train_losses = []
    val_losses = []
    for epoch in range(1,epochs+1):
        all_preds = []
        all_labels = []
        loss_train = 0.0
        correct_train = 0
        total =0
        total_batch =0
        model.train()
        for batch_id, data in enumerate(tqdm(train_loader)):
            
            data = data.to(device)
            
            out = model(data)
            with torch.no_grad():
                pred = out.argmax(dim=1)
                
            loss = loss_fn(out, data['frame_window'].y) 
            optimizer.zero_grad()
            loss.backward()        
            optimizer.step() 
            loss_train  += loss.item()  
            correct_train += int((pred == data['frame_window'].y).sum())
            total += len(data)
            total_batch += 1
        model.eval()    
        total_val, total_val_batch, loss_val, correct_val, all_labels_epoch,all_preds_epoch = validate_model()
        scheduler.step(loss_val/total_val)
        current_lr = get_lr()
        all_labels.extend(all_labels_epoch)
        all_preds.extend(all_preds_epoch) 
        train_losses.append(loss_train / total_batch)
        val_losses.append(loss_val / total_val_batch)
        print('{} Epoch-{}, train-loss {:.4f}, train-acc {:.4f} || val-loss {:.4f}, val-acc {:.4f} || lr - {}'.format(
            datetime.datetime.now(), epoch, loss_train/total_batch, correct_train/total, loss_val/total_val_batch,
            correct_val/total_val, current_lr))
        # Check to see if there is early stopping
        if early_stopper.early_stop(loss_val):
            break
            
    plt.plot(range(1, epoch + 1), train_losses, label='Train Loss')
    plt.plot(range(1, epoch + 1), val_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()
    
    cm = confusion_matrix(all_labels, all_preds)
    cmn = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    overall_accuracy = accuracy_score(all_labels, all_preds)
    print(overall_accuracy)
    class_accuracy = cm.diagonal() / cm.sum(axis=1)
    for class_idx, acc in enumerate(class_accuracy):
        print(f'Class {class_idx} Accuracy: {acc}')
    # Plot confusion matrix
    plt.figure(figsize=(15, 6))
    plt.subplot(1, 2, 1)
    sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', xticklabels=range(len(class_accuracy)), yticklabels=range(len(class_accuracy)))
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.show()
    
    plt.subplot(1, 2, 2)
    sns.heatmap(cmn, annot=True, fmt='.2%', cmap='Blues', xticklabels=range(len(class_accuracy)), yticklabels=range(len(class_accuracy)))
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Normalized Confusion Matrix')
    plt.show()
        
def validate_model():
    all_preds_epoch = []
    all_labels_epoch = []
    model.eval()

    correct = 0
    total = 0
    total_batch = 0
    loss_val =0.0

    
    with torch.no_grad():
        for batch_id, data in enumerate(tqdm(val_loader)):
            data = data.to(device)
            out = model(data)
            pred = out.argmax(dim=1)
            total += len(data)
            total_batch += 1
            correct += int((pred == data['frame_window'].y).sum())
            loss = loss_fn(out, data['frame_window'].y)
            loss_val += loss.item()
            all_preds_epoch.extend(pred.cpu().numpy())
            all_labels_epoch.extend(data['frame_window'].y.cpu().numpy()) 
    return total, total_batch, loss_val, correct,all_labels_epoch,all_preds_epoch  

class EarlyStopper:
    """
    The EarlyStopper class for training
    """

    def __init__(self, patience=30, min_delta=0.015):
        """
        Initialization

        :param patience: Number of epochs to wait
        :param min_delta: The minimum change that is required
        """

        # Init
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.min_validation_loss = np.inf

    def early_stop(self, validation_loss):
        """
        Early stop instance

        :param validation_loss: The loss to monitor
        :return:
        """

        if validation_loss < self.min_validation_loss:
            self.min_validation_loss = validation_loss
            self.counter = 0
        elif validation_loss > (self.min_validation_loss + self.min_delta):
            self.counter += 1
            if self.counter >= self.patience:
                sys.stdout.write("EarlyStopping the training process due to non changing validation loss\n")
                return True
        return False
    
def get_lr():

    """
    Get the current learning rate for the model
    :return: The learning rate for the current epoch after the update
    """

    for param_group in optimizer.param_groups:
        return param_group["lr"]

Thanks in advance!