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.