File size: 6,082 Bytes
3c93691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | # Model Card: NextViT-Small for Facial Expression Recognition (FER)
This model card provides comprehensive details, training configurations, and performance benchmarks for the fine-tuned **NextViT-Small** model adapted for Facial Expression Recognition (FER).
## Model Details
* **Model Name:** NextViT-Small FER
* **Architecture:** NextViT-Small (`nextvit_small`), a hybrid CNN-Transformer vision model designed for high-performance and hardware-efficient inference on edge devices.
* **Developer:** Fine-tuned locally.
* **Task:** Multi-class image classification for Facial Expression Recognition (7 classes).
* **Number of Classes:** 7 (`angry`, `disgust`, `fear`, `happy`, `neutral`, `sad`, `surprise`).
* **Parameters:** ~28 million parameters (~112 MB for float32 weights-only; checkpoint size is ~381.75 MB including optimizer state).
* **Base Model Weight Source:** Pre-trained on ImageNet (`ckpt5.pth`).
## Training Configuration & Hyperparameters
The model was fine-tuned for 50 epochs on a custom facial expression dataset using the following training parameters:
| Hyperparameter | Value | Description |
| :--- | :--- | :--- |
| **Optimizer** | AdamW | Weight decay set to `0.01` |
| **Learning Rate (Base)** | `2e-4` | Cosine learning rate scheduler with warmup |
| **Epochs** | `50` | Total training epochs |
| **Batch Size** | `128` | Per-GPU batch size (effective batch size of 256 on 2 GPUs) |
| **Input Image Size** | `224 x 224` | Bilinear resize and normalized to ImageNet statistics |
| **Data Augmentations** | RandAugment (`rand-m5-mstd0.5-inc2`), Random Erasing (`reprob 0.1`) | Mixup and Cutmix were disabled (`mixup 0`, `cutmix 0`) |
| **Distributed Training** | PyTorch Distributed Data Parallel (DDP) | 2 GPUs (Distributed evaluation enabled, NCCL backend) |
| **Total Training Time** | `0:41:41` (41m 41s) | Training completed on a dual-GPU node |
---
## Evaluation Metrics & Performance
### 1. General Test Set Performance (12,957 Images)
At the end of training (Epoch 50), the model achieved the following performance on the full validation/test dataset:
* **Top-1 Accuracy:** **82.20%**
* **Top-5 Accuracy:** **99.04%**
* **Final Test Loss:** **0.739**
### 2. Detailed Test Split Performance (2,463 Images)
When evaluated on a dedicated test split of 2,463 samples, NextViT-Small achieved:
* **Overall Accuracy:** **83.88%**
* **Precision (Macro):** **75.92%**
* **Precision (Weighted):** **83.82%**
* **Recall (Macro):** **73.93%**
* **Recall (Weighted):** **83.88%**
* **F1-Score (Macro):** **74.61%**
* **F1-Score (Weighted):** **83.78%**
#### Per-Class Performance Breakdown:
| Class | Class ID | Precision | Recall | F1-Score | Support |
| :--- | :---: | :---: | :---: | :---: | :---: |
| **happy** | 3 | 95.51% | 93.37% | 94.43% | 935 |
| **surprise** | 6 | 83.33% | 84.34% | 83.83% | 249 |
| **angry** | 0 | 79.38% | 83.01% | 81.15% | 153 |
| **neutral** | 4 | 77.84% | 81.89% | 79.81% | 519 |
| **sad** | 5 | 78.67% | 80.39% | 79.52% | 413 |
| **disgust** | 1 | 58.78% | 56.62% | 57.68% | 136 |
| **fear** | 2 | 57.89% | 37.93% | 45.83% | 58 |
---
### 3. Comparison with ViT-Small Baseline
NextViT-Small shows significant improvements over a standard vision transformer (ViT-Small) baseline evaluated on the same 2,463 image test split:
* **ViT-Small Accuracy:** **80.31%** (Macro F1-Score: **71.39%**)
* **NextViT-Small Accuracy:** **83.88%** (Macro F1-Score: **74.61%**)
* **Performance Gain:** **+3.57% Accuracy** improvement with NextViT-Small.
NextViT-Small demonstrates much stronger capability on classes with lower support (e.g. `angry`, `disgust`) due to its hybrid convolutional-attention structure, which excels in learning fine-grained local facial details.
---
## How to Use
To run inference using the trained NextViT-Small model, follow this example snippet:
```python
import os
import sys
import torch
from PIL import Image
from torchvision import transforms
# 1. Path setup: Add Next-ViT classification folder to path so timm can register nextvit_small
NEXTVIT_PATH = "/path/to/Next-ViT/classification" # Update as appropriate
sys.path.append(NEXTVIT_PATH)
import nextvit # registers 'nextvit_small' with timm
import timm
# 2. Load the model architecture and pretrained weights
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = timm.create_model("nextvit_small", num_classes=7)
checkpoint_path = "checkpoint_best.pth" # Update path to best checkpoint
checkpoint = torch.load(checkpoint_path, map_location=device)
if "model" in checkpoint:
model.load_state_dict(checkpoint["model"])
else:
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
# 3. Define the image preprocessing transforms
image_processor = 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)),
])
# 4. Run inference
class_names = ["angry", "disgust", "fear", "happy", "neutral", "sad", "surprise"]
def predict_expression(image_path):
img = Image.open(image_path).convert("RGB")
tensor = image_processor(img).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(tensor)
probs = torch.nn.functional.softmax(outputs, dim=-1)
pred_idx = torch.argmax(probs, dim=-1).item()
return class_names[pred_idx], probs[0][pred_idx].item()
# Example Prediction
# expression, confidence = predict_expression("test_face.jpg")
# print(f"Predicted Expression: {expression} ({confidence:.2%})")
```
---
## Intended Use & Limitations
* **Intended Use:** Facial Expression Recognition in user-facing applications (e.g. human-computer interaction, affective computing, user experience testing).
* **Limitations:** Performance on expressions with very low support (such as `fear` or `disgust`) remains relatively lower compared to well-represented expressions (like `happy` and `neutral`). Users should be cautious when deploying in scenarios sensitive to false negatives on these classes.
|