| # 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. | |
|  | |