harveymannering/VGG_AffectNet7

VGG Finetuned on AffectNet

VGG model taken finetuned on AffectNet data for prediction of the 7 basic emotions. The model architecture can be described in code as follows:

from torchvision import models
import torch.nn as nn
import torch
from huggingface_hub import hf_hub_download

class CustomVGG(nn.Module):
    def __init__(self):
        super(CustomVGG, self).__init__()
        
        # Download VGG model 
        self.vgg = models.vgg16(pretrained=True)
        
        # Add a final MLP to be run after the VGG model
        self.vgg.classifier[6] = nn.Linear(in_features=4096, out_features=7)

    def forward(self, x):
        # Get features up to classifier[4] (second-to-last layer)
        features = self.vgg.features(x)
        features = self.vgg.avgpool(features)
        features = torch.flatten(features, 1)
        
        # Pass through first 5 classifier layers
        for i in range(5):
            features = self.vgg.classifier[i](features)
        
        second_to_last = features  # Features before final layer
        pred = self.vgg.classifier[6](features)  # Final prediction
        
        return pred, second_to_last


# Load model
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weights_path = hf_hub_download(repo_id="harveymannering/VGG_AffectNet7", filename="vgg_7_best.pth")
classifier = CustomVGG()
classifier.load_state_dict(torch.load(weights_path, map_location=device))
classifier = classifier.to(device)

Here is the loss plot for this training run. This checkpoint is taken from the epoch with the best validation loss. At it's peak it acheievd 59% accuracy on the validation set. It was trained on 7 basic emotion classes with no face cropping, but a bunch of augmentations including the standard ImageNet normalization.