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