Update README.md

README.md

@@ -1 +1,25 @@
-VGG model taken finetuned on AffectNet data for prediction of the 7 basic emotions.
+VGG model taken finetuned on AffectNet data for prediction of the 7 basic emotions. The model architecture can be described in code as follows: 
+
+```
+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):
+        pred = self.vgg(x)
+        return pred
+    
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+classifier = CustomVGG()
+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.
+
+![480942868-7b13906a-58a0-428a-bd16-b359eb89301d](https://cdn-uploads.huggingface.co/production/uploads/6349716695ab8cce385f450e/cOtOMddRH0IIzBTd4bgb0.png)