peepeeyanto/ViTFERPP_vit_small_patch16_224.augreg_in1k

ViT-Small for Facial Expression Recognition (FER++)

This repository contains a Vision Transformer (ViT-Small) model fine-tuned for Facial Expression Recognition (FER).

The model is initialized from the pretrained checkpoint vit_small_patch16_224.augreg_in1k from the timm library and fine-tuned on a 7-class emotion classification dataset (aligned with FER++ taxonomy).

Model Details

• Model Architecture: Vision Transformer (ViT-Small)
• Pretrained Base: vit_small_patch16_224.augreg_in1k
• Number of Parameters: ~22M
• Task: 7-class Facial Expression Classification
• Classes:
  • 0: Angry
  • 1: Disgust
  • 2: Fear
  • 3: Happy
  • 4: Neutral
  • 5: Sad
  • 6: Surprise

Training Parameters & Hyperparameters

The model was fine-tuned using PyTorch and the timm library with the following setup:

Hyperparameter	Value
Epochs	50
Batch Size	128
Base Learning Rate	5e-4
Weight Decay	0.01
Optimizer	AdamW
Scheduler	Cosine Annealing (with 5-epoch warmup starting at 1e-6, min LR 1e-5)
Image Resolution	224 x 224 (Bicubic Interpolation)
Stochastic Depth (Drop Path)	0.1
Precision	Mixed Precision (AMP)
Data Augmentations	RandAugment, Random Erasing (p=0.1)
Loss Function	Weighted Cross-Entropy Loss
Class Weights	[1.560, 3.737, 2.242, 0.541, 0.527, 0.970, 1.149] (mapping to classes 0 to 6)

---

Evaluation Results

Validation Set Performance

• Accuracy: 83.83%
• Precision (Macro): 0.7892
• Recall (Macro): 0.7767
• F1-Score (Macro): 0.7817
• F1-Score (Weighted): 0.8378

Per-Class Validation Metrics:

Class ID	Class Name	Precision	Recall	F1-Score	Support
3	Happy	0.9247	0.9345	0.9295	473
6	Surprise	0.8636	0.7661	0.8120	124
4	Neutral	0.7943	0.8145	0.8043	275
0	Angry	0.7412	0.8400	0.7875	75
5	Sad	0.7758	0.7711	0.7734	166
2	Fear	0.8000	0.7273	0.7619	33
1	Disgust	0.6250	0.5833	0.6034	60

---

Test Set Performance

• Accuracy: 83.68%
• Precision (Macro): 0.7800
• Recall (Macro): 0.7413
• F1-Score (Macro): 0.7565
• F1-Score (Weighted): 0.8348

Per-Class Test Metrics:

Class ID	Class Name	Precision	Recall	F1-Score	Support
3	Happy	0.9187	0.9378	0.9281	482
6	Surprise	0.8793	0.7969	0.8361	128
0	Angry	0.8158	0.8493	0.8322	73
5	Sad	0.8688	0.7277	0.7920	191
4	Neutral	0.7240	0.8707	0.7906	232
2	Fear	0.6471	0.5000	0.5641	22
1	Disgust	0.6066	0.5068	0.5522	73

---

How to Use

Here is how you can load the model and run inference on an image using PyTorch and timm:

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from timm import create_model

# Define class mapping
class_mapping = {
    0: "Angry",
    1: "Disgust",
    2: "Fear",
    3: "Happy",
    4: "Neutral",
    5: "Sad",
    6: "Surprise"
}

# 1. Define image preprocessing matching validation/test set config
transform = transforms.Compose([
    transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

# 2. Instantiate and build the model structure
model = create_model('vit_small_patch16_224.augreg_in1k', pretrained=False, num_classes=7)

# 3. Load weight state dictionary
state_dict = torch.load('vitsmall2.pth', map_location='cpu')

# Clean state_dict if it was saved using DataParallel (remove 'module.' prefix)
if list(state_dict.keys())[0].startswith('module.'):
    state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

model.load_state_dict(state_dict)
model.eval()

# 4. Perform Inference
image_path = "path_to_facial_image.jpg"
image = Image.open(image_path).convert('RGB')
input_tensor = transform(image).unsqueeze(0)  # Add batch dimension

with torch.no_grad():
    outputs = model(input_tensor)
    probabilities = torch.softmax(outputs, dim=1)[0]
    predicted_class_id = torch.argmax(probabilities).item()

print(f"Predicted emotion: {class_mapping[predicted_class_id]} ({probabilities[predicted_class_id].item():.2%})")