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