Clustering — Face Recognition (PyTorch)

PyTorch clustering unsupervised embeddings face recognition metric learning

Example: weightslab/examples/PyTorch/wl-clustering/main.py

Task: Metric learning with triplet loss on the Olivetti / LFW face dataset. The goal is to train an embedding network so that embeddings from the same person cluster together.

This example shows WeightsLab used in a contrastive / metric-learning setting where there is no standard per-sample label — the signal of interest is the embedding distance.

Integration walkthrough

1. Register hyperparameters and data

wl.watch_or_edit(parameters, flag="hyperparameters",
                 defaults=parameters, poll_interval=1.0)

train_loader = wl.watch_or_edit(
    train_dataset,
    flag="data",
    loader_name="train_loader",
    batch_size=parameters["batch_size"],
    shuffle=True,
    compute_hash=False,
    num_workers=parameters.get("num_workers", 0),
)

The dataset yields image triplets (anchor, positive, negative) plus their stable UIDs. WeightsLab records which triplets the model has seen and lets you inspect the hardest negatives in the studio.

2. Guard contexts — same as classification

def train(loader, model, optimizer, device):
    with guard_training_context:
        images, uids, labels, _ = next(loader)
        embeddings = model(images.to(device))
        triplet_loss = compute_triplet_loss(embeddings, labels)
        triplet_loss.mean().backward()
        optimizer.step()

def evaluate(loader, model, device):
    with guard_testing_context, torch.no_grad():
        for images, uids, labels, _ in loader:
            embeddings = model(images.to(device))
            ...

3. Saving embedding metrics per sample

metrics = FaceMetrics.compute_all_metrics(embeddings, labels, uids)
wl.save_signals(
    {
        "train/intra_class_dist":  metrics["intra_class_dist"],
        "train/inter_class_dist":  metrics["inter_class_dist"],
        "train/silhouette_score":  metrics["silhouette_score"],
    },
    batch_ids=uids,
    origin="train_loader",
)

Because there is no single “loss per sample” in metric learning, you compute whichever distances are meaningful and push them through wl.save_signals directly. The studio plots these curves and lets you sort images by silhouette score to identify consistently confused identities.

Tip

This example is bundled with WeightsLab:

weightslab start                # 1. deploy the studio
weightslab start example --clus # 2. start the clustering demo