Wafer-Map Defect Classification

A convolutional neural network that classifies spatial defect patterns on semiconductor wafer maps, trained on the WM-811K dataset. Built to run natively on Apple Silicon (M1/M2) via PyTorch's MPS backend, with free tools end to end.

Problem

Wafer maps encode which dies on a silicon wafer passed or failed testing. The spatial pattern of failures (a ring at the edge, a central blob, a scratch) points to different root causes in the fab. This model classifies a map into one of 9 categories: 8 defect patterns (Center, Donut, Edge-Loc, Edge-Ring, Loc, Random, Scratch, Near-full) plus "none".

The core challenge is class imbalance — "none" dominates and several defect classes are rare — so the project leans on class weighting, oversampling, augmentation, and macro-F1 rather than raw accuracy.

Setup

Requires uv.

git clone <your-repo-url>
cd wafer-defect-cnn
uv sync

Get the data

The dataset is ~2GB and is not committed. Download it from Kaggle:

# Needs a Kaggle API token at ~/.kaggle/kaggle.json
uv run kaggle datasets download -d qingyi/wm811k-wafer-map
unzip wm811k-wafer-map.zip -d data/

Usage

uv run python -m src.train

Auto-selects the MPS GPU on Apple Silicon (falls back to CPU elsewhere), then trains every architecture registered in src/ml_model_registry.py back to back on identical stratified splits, printing per-epoch validation macro-F1 and a final ranked comparison. Each architecture gets its own checkpoint (checkpoints/<name>.pt) and confusion-matrix PNG (outputs/confusion_matrix_<name>.png) so runs don't clobber each other.

To train/compare a subset instead of all 4, call compare([...]) from src/train.py directly, e.g.:

uv run python -c "from src.train import compare; compare(['res_block', 'resnet18'])"

Pass --exclude-none to drop the "none" class from the data entirely and retrain with an 8-way output head, instead of the default 9-way run:

uv run python -m src.train --exclude-none

Project layout

.
├── src/
│   ├── config.py               # paths, hyperparameters, classes, device
│   ├── DataHandler/
│   │   ├── loader.py           # pandas shim, pickle loading/caching
│   │   ├── preprocessing.py    # to_fixed_map (resize + recode)
│   │   └── dataset.py          # WaferDataset
│   ├── MachineLearningModels/
│   │   ├── ThreeBlockCNN.py
│   │   ├── ResBlockCNN.py
│   │   ├── SqueezeAndExtractCNN.py
│   │   └── ResNet18Transfer.py
│   ├── ml_model_registry.py    # build_model() factory over the 4 architectures
│   ├── sampling.py             # class weights + oversampling for imbalance
│   ├── evaluation.py           # macro-F1 + confusion matrix
│   ├── visualization.py        # confusion-matrix plotting
│   ├── gradcam.py              # Grad-CAM: what each model attends to
│   └── train.py                # splits, dataloaders, training loop, comparison runner
├── notebooks/
│   ├── data_exploration.ipynb     # class distribution, example maps per class, pixel sanity check
│   └── gradcam_exploration.ipynb  # Grad-CAM per class, and across all 4 architectures
├── BUILD_PLAN.md     # detailed week-by-week build plan
├── pyproject.toml    # dependencies (uv)
├── .gitignore
└── data/             # WM-811K pickle (not committed)

Results

Test macro-F1 on the natural (imbalanced) 9-class distribution, stratified split:

Model	Test macro-F1
resnet18	0.706
res_block	0.674
se_cnn	0.632
three_block	0.532

resnet18's pretrained ImageNet features generalize best despite the domain mismatch with wafer maps, though Grad-CAM shows its explanations are much coarser than the custom CNNs' (see gradcam_exploration.ipynb) — its last feature map is only ~2x2 before upsampling, vs. 8x8 for the other three.

8-class-only (no "none")

Dropping "none" from the data entirely — not just filtering it out of eval — and retraining with an 8-way output head, for a fair comparison against the reference SVM baseline:

uv run python -m src.train --exclude-none

Model	Test macro-F1 (8-class)
resnet18	0.898
res_block	0.854
se_cnn	0.709
three_block	0.609

Same ranking as the 9-class run, but every model scores much higher once "none" (the overwhelming majority class) is out of the picture — macro-F1 was being dragged down by confusion between rare defect classes and "none", not by confusion among the defect classes themselves. Checkpoints/plots use an _8class suffix (checkpoints/<name>_8class.pt, outputs/confusion_matrix_<name>_8class.png) so they don't overwrite the 9-class artifacts.

Roadmap

• Visualization script (example maps per class + class distribution) — notebooks/data_exploration.ipynb
• Grad-CAM heatmaps to see what the CNN attends to — src/gradcam.py, notebooks/gradcam_exploration.ipynb
• Compare small CNN vs. fine-tuned ResNet18 — see Results above
• True 8-class-only retrain (dropping "none" entirely, not just filtering eval) — see 8-class-only results above
• Multi-label head for maps with mixed defect patterns — blocked on data: WM-811K only has single-label annotations, would need a mixed-defect dataset (e.g. MixedWM38) or synthetic overlays
• Ablation: to_fixed_map's mode="bilinear" resize vs. mode="area"/"nearest"

License

MIT