Fingerprinting ViT-GMM

Unsupervised and Semi-Supervised Source Fingerprinting of Additive-Manufactured Parts Using Vision Transformers and Gaussian Mixture Models

---

Overview

This repository contains the code for our paper on source fingerprinting of additively manufactured (AM) connectors and plugs. We train a Vision Transformer (ViT-MAE) backbone jointly with a Gaussian Mixture Model (GMM) clustering head to identify the printer, build location, material, and machine generation from macro-scale images — without requiring labeled data during training.

Key features:

• Fully unsupervised fingerprinting via GMM negative log-likelihood loss
• Semi-supervised extension using a fraction of class labels
• Cross-domain generalization from connectors to plugs
• Multi-GPU distributed training with PyTorch DDP + SLURM

---

Requirements

pip install -r requirements.txt

GPU-accelerated libraries (optional, for t-SNE visualization):

Install via RAPIDS:

conda install -c rapidsai -c conda-forge cuml cupy

Tested environment

• Python 3.10+
• PyTorch 2.1 + CUDA 12.x
• 4× H200 80 GB GPUs (training), 8× L40S 48 GB GPUs (smoke test)

---

Model Weights

Pre-trained ViT-MAE weights are hosted on HuggingFace Hub:

from huggingface_hub import hf_hub_download

# ViT-Small encoder (default)
hf_hub_download(repo_id="wpklab/Fingerprinting-ViT-GMM", filename="ae_best_encoder_2.pth",
                local_dir="data/Models")

# ViT-Large encoder (small_ae: false)
hf_hub_download(repo_id="wpklab/Fingerprinting-ViT-GMM",
                filename="printer_identification_full_ae.pth", local_dir="data/Models")

Note: The HuggingFace repo will be made public upon paper acceptance.

---

Dataset

Datasets will be released at publication. Expected structure:

data/
├── printer_DLS_connector_11/          # 11-class printer ID — connector images
│   ├── train/{M2-1,M2-2,...}/
│   └── val/{M2-1,M2-2,...}/
├── printer_DLS_plug_11/               # 11-class printer ID — plug images
├── buildloc_DLS_connector_36_48/      # 36-class build location
├── build_DLS_connector_49_36/         # 49-class build number
├── material_DLS_connector_3_356/      # 3-class material (RPU, EPX, UMA)
└── printergen_DLS_connector_2_250/    # 2-class machine generation (M2, L1)

Datasets will be available on:

• Kaggle Datasets — TBD
• Zenodo — TBD (with DOI for citation)

---

Running Experiments

Standalone (single or multi-GPU, no sweep)

# Printer source identification — fully unsupervised
python run.py --config configs/printer_detection.yaml

# Cross-domain: train on connectors, evaluate on plugs
python run.py --config configs/connector_plug.yaml

# Build location classification
python run.py --config configs/build_location.yaml

# Connector attribute classification (material or machine generation)
python run.py --config configs/connector_attributes.yaml

# Anomaly detection (hold out one printer class)
python run.py --config configs/printer_anomaly.yaml

W&B Hyperparameter Sweep

# Create a sweep on W&B first, then launch an agent:
python run.py --config configs/printer_detection.yaml \
              --sweep_id <entity>/<project>/<sweep_id>

SLURM (multi-GPU cluster)

# Single experiment
sbatch slurm/train.sbatch configs/printer_detection.yaml

# With a W&B sweep
sbatch slurm/train.sbatch configs/printer_detection.yaml <entity>/<project>/<sweep_id>

Edit slurm/train.sbatch to set your --account and --partition before submitting.

---

Configuration

All hyperparameters are documented in configs/base.yaml. Experiment configs override only the fields that differ from the base:

Config	Dataset	Task
printer_detection.yaml	connector_connector	11-class printer ID
printer_anomaly.yaml	connector_connector	One-class anomaly detection
connector_attributes.yaml	material / machine_generation	Attribute classification
build_location.yaml	build_loc	36-class build location
connector_plug.yaml	connector_plug	Cross-domain generalization

---

Repository Structure

Fingerprinting-ViT-GMM/
├── main.py                    # SLURM distributed entry point
├── run.py                     # Local / single-process entry point
├── requirements.txt
├── configs/
│   ├── base.yaml              # All hyperparameter defaults (documented)
│   ├── printer_detection.yaml
│   ├── printer_anomaly.yaml
│   ├── connector_attributes.yaml
│   ├── build_location.yaml
│   └── connector_plug.yaml
├── slurm/
│   ├── train.sbatch           # Multi-GPU H200 job script
│   └── smoke_test.sbatch      # Quick correctness check (L40S)
├── tests/
│   └── make_smoke_data.py     # Generate synthetic datasets for smoke testing
├── src/
│   ├── trainer.py             # Core training loop (DDP, GMM, semi-supervised)
│   ├── datasets.py            # SubsetInMemoryDataset, SourceTrackedDataset
│   ├── models.py              # FeatureReducer, ClusteringLayer, ResNetAutoencoder
│   ├── losses.py              # GMM, entropy, separation, and SwAV losses
│   ├── math_utils.py          # PyTorch GMM, PCA, majority voting
│   ├── transforms.py          # PadToSize, random crops
│   └── utils.py               # EarlyStopping, clustering accuracy
└── assets/
    └── model_overview.png     # Paper figure

---

Citation

@article{,
  title   = {},
  author  = {},
  journal = {},
  year    = {2025},
}

---

License

This project is released under the MIT License. See LICENSE for details.