Update README.md

README.md

@@ -1,6 +1,13 @@
+# 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__()
@@ -12,12 +19,28 @@ class CustomVGG(nn.Module):
         self.vgg.classifier[6] = nn.Linear(in_features=4096, out_features=7)
 
     def forward(self, x):
-        pred = self.vgg(x)
-        return pred
-    
+        # 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.