ViT Defect Classification (Small-Data Experiments)

This project fine‑tunes a Vision Transformer (ViT) model for binary defect detection (defect vs normal) on a very small dataset (~190 images), and systematically compares different tuning strategies in a low‑data regime.

The goal is not to reach SOTA accuracy, but to understand:

• Which tuning strategy is more robust on small data.
• Whether the model has really learned a meaningful representation (latent space & attention), instead of just memorizing the training set.

1. Project Structure

---

.
├─ data_dir/
│  ├─ train/
│  │  ├─ defect/
│  │  └─ normal/
│  ├─ validation/   # or "valid"/"val"
│  │  ├─ defect/
│  │  └─ normal/
│  └─ test/
│     ├─ defect/
│     └─ normal/
├─ train_vit.ipynb        # baseline training + plots
├─ experiment.ipynb       # small-data experiments (3 strategies, t-SNE, attention)
├─ test.ipynb             # inference & test evaluation
├─ vit-output/            # baseline fine-tuned checkpoint
└─ vit-experiments/
   ├─ linear-probe/
   ├─ lora/
   └─ full-ft/

2. Environment & Dependencies

---

Tested with:

• Python 3.10–3.13
• PyTorch 2.10.0, torchvision 0.25.0
• transformers==4.45.0, accelerate>=1.1.0
• datasets, evaluate, matplotlib, pandas, pillow
• scikit-learn, umap-learn, peft

Install (recommended in a virtual environment):

pip install datasets evaluate torch torchvision matplotlib pandas \
  "transformers==4.45.0" "accelerate>=1.1.0" tf-keras scikit-learn umap-learn peft

3. Dataset Layout

---

Expected layout under data_dir:

data_dir/
├─ train/
│  ├─ defect/
│  └─ normal/
├─ validation/   # or "valid"/"val"
│  ├─ defect/
│  └─ normal/
└─ test/
   ├─ defect/
   └─ normal/

Each subfolder contains images of that class.
Update the path in the notebooks if needed:

data_dir: str = r"C:\Users\user\Documents\ViT\data_dir"

4. Baseline Training (train_vit.ipynb)

---

train_vit.ipynb provides a simple baseline:

• Loads google/vit-base-patch16-224-in21k.
• Fine‑tunes on the small dataset.
• Evaluates on validation set.
• Saves:
  • fine‑tuned model to ./vit-output
  • training logs and plots to ./vit-output/plots/:
    • training_loss.png
    • epoch_loss.png
    • epoch_accuracy.png
    • learning_rate.png
    • training_logs.csv

Use this notebook as the minimal, “clean” reference implementation.

5. Small-Data ViT Experiment (experiment.ipynb)

---

5.1 Experimental Setup

- **Base model**: `google/vit-base-patch16-224-in21k`
- **Dataset**: ~150 train, 30 validation, 20 test images
- **Task**: binary classification (`defect` vs `normal`)
- **Frameworks**: Hugging Face `transformers`, `datasets`, `peft` (for LoRA), `sklearn`, `umap-learn`

All three experiments share:

- Same train/validation split
- Same optimizer settings (learning rate, batch size, epochs = 5)
- Same data preprocessing (ViT image processor, 224×224 resize, normalization)

5.2 Tuning Strategies Compared

5.2.1 Linear Probing ^^^^^^^^^^^^^^^^^^^^

• What: Freeze the entire ViT backbone, only train the final classification head.
• How:
  • param.requires_grad = False for all layers
  • param.requires_grad = True only for model.classifier
• Why:
  • Uses the pretrained representation as‑is.
  • Very few trainable parameters → low risk of overfitting on small data.

5.2.2 LoRA Fine-tuning ^^^^^^^^^^^^^^^^^^^^^^

• What: Apply Low‑Rank Adaptation (LoRA) to the attention layers, using the peft library.
• How (example):
  • Target modules: ["query", "key", "value"]
  • Rank r = 8, dropout 0.05
  • Most original weights are frozen; only a small number of adapter parameters are trained (≈1% of weights).
• Why:
  • Current (2026) popular method for adapting large models on small datasets.
  • Compromise between flexibility and regularization.
• Note:
  • In the first implementation the validation metrics for LoRA were not logged correctly (only runtime was reported), so the quantitative result for LoRA is inconclusive. This is a limitation of the current experiment, not of LoRA itself.

5.2.3 Full Fine-tuning ^^^^^^^^^^^^^^^^^^^^^^

• What: Unfreeze all layers and train the entire model.
• Why:
  • Maximum flexibility, but highest risk of overfitting when there are only ~190 images.

6. Metrics and Results

---

• Validation metric: eval_accuracy (main) and eval_loss.
• For each strategy, the notebook records:
  • eval_accuracy
  • eval_loss
  • Training curves under vit-experiments/<strategy>/logs and related plots.

6.1 Summary Table (current run)

=================  ==========  ==============
Strategy           Eval loss   Eval accuracy
=================  ==========  ==============
Linear probe       0.6674      0.6333
LoRA               —           —
Full fine‑tune     0.0404      1.0000
=================  ==========  ==============

**Observations:**

- **Linear probe** reaches around **63% validation accuracy**, only slightly better than random (50%). This suggests that **just training the head is not enough** to fully adapt ViT on this dataset.
- **Full fine‑tuning** quickly reaches **100% validation accuracy with very low loss** on only ~30 validation images. Combined with the small data size, this is strong evidence of **overfitting**: the model can easily memorize the small dataset.
- **LoRA** in this initial run does not report accuracy due to a logging/configuration issue, so its performance is **not yet measured**. Fixing the metric logging and re‑running LoRA is an important next step.

These results illustrate a typical small‑data trade‑off:

- Full fine‑tuning gives very strong numbers on the tiny validation set, but may not generalize.
- Linear probing is more conservative and underfits.
- Parameter‑efficient methods like LoRA are expected to lie **between** these two, but require a correct evaluation setup.

7. Latent Space Visualization (t‑SNE / UMAP)
--------------------------------------------

To understand what the model has learned beyond scalar accuracy, `experiment.ipynb`:

1. **Extracts CLS embeddings**  
   - For each image (train + validation + test), takes the 768‑dimensional `[CLS]` token from the last ViT layer.
2. **Performs dimensionality reduction**  
   - Applies **t‑SNE** and/or **UMAP** to project 768‑D vectors into 2‑D.
3. **Plots**  
   - Colors points by class: `defect` vs `normal`.

**Research question**:

- If the two classes form **well‑separated clusters** in 2‑D, this is evidence that the model has learned a meaningful representation, even with a small dataset.
- If the clusters are highly mixed, it suggests that the current fine‑tuning strategy or data size is not sufficient.

You can generate one figure per strategy (linear probe / LoRA / full FT) to visually compare how “clean” the separation is, and also to check whether the apparently perfect full‑FT accuracy corresponds to a clearly separated latent space or just memorization.

8. Attention Map Interpretability
---------------------------------

To check **where** the ViT is looking:

1. For a selected image, run the model with `output_attentions=True` and `attn_implementation="eager"` (required for `transformers>=4.41` to actually return attention weights for ViT).
2. Take the last layer’s attention from the **CLS token to all patch tokens**.
3. Reshape this into a 2‑D grid and upsample to the original image size.
4. Overlay the attention heatmap on top of the image.

**Interpretation**:

- If high attention weights concentrate on the **defect region** (scratch, crack, etc.), this supports that the model is using the correct cues.
- If attention focuses on irrelevant background (corners, borders, labels), the model might be “cheating”, and the current setup is not reliable.

9. How to Run
-------------

9.1 Baseline
~~~~~~~~~~~~

1. Open `train_vit.ipynb`.
2. Run all cells.
3. Check:
   - Validation metrics in the last cell.
   - Plots under `vit-output/plots/`.

9.2 Experiments
~~~~~~~~~~~~~~~

1. Open `experiment.ipynb`.
2. Run:
   - Setup cells (imports, config, data loading).
   - The three `run_experiment` calls (linear probe, LoRA, full FT).
   - The comparison table cell.
   - Latent‑space cells (t‑SNE / UMAP).
   - Attention map cells.

10. What This Experiment Shows
------------------------------

This small experiment is designed to:

- Demonstrate the **risk of full fine‑tuning** on small data (very high validation accuracy on very few samples).
- Show that **simple linear probing underfits**, indicating the head alone cannot fully exploit the pretrained features for this task.
- Motivate **parameter‑efficient methods (such as LoRA)** as a promising direction, while also highlighting that the current LoRA run still needs correct metric logging.

Even though the dataset is small, combining:

- quantitative metrics (validation accuracy / loss), and  
- qualitative analysis (latent space, attention)

gives a convincing story about how different fine‑tuning strategies behave in a low‑data regime and why careful evaluation is necessary.