Generation / Anomaly Detection — MVTec (PyTorch)

PyTorch anomaly detection generation unsupervised mvtec reconstruction

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

Task: Unsupervised anomaly detection on MVTec capsule images with a multi-task UNet (classification head + reconstruction head + contrastive loss).

This example introduces two advanced patterns:

• compute_dependencies=False when wrapping a complex generative model.

• wl.save_group_signals for signals computed over pairs of samples.

Integration walkthrough

1. Multi-task model wrapping

model = wl.watch_or_edit(
    _model,
    flag="model",
    device=device,
    compute_dependencies=False,
)

compute_dependencies=False skips the static dependency graph computation for the wrapped model. Use this when the model has dynamic control flow, multiple outputs, or cannot be traced by torch.fx — common in encoder-decoder architectures.

2. Dataset with paired samples

The dataset yields two images per item (an anchor and a randomly paired image) so the contrastive loss can be computed in-batch:

train_loader = wl.watch_or_edit(
    train_dataset,
    flag="data",
    loader_name="train_loader",
    batch_size=parameters["batch_size"],
    shuffle=True,
    compute_hash=False,
)
# Loader yields: [img1, img2], [uid1, uid2], [label1, label2], metadata

The dataset’s __getitem__ returns a pair so each WeightsLab batch contains two groups of sample IDs.

3. Per-sample signals for individual losses

with guard_training_context:
    [img1, img2], [uid1, uid2], [label1, label2], _ = next(train_loader)
    cls_out, recon_out = model([img1, img2])

    cls_loss_per  = wl.save_signals({"train/cls_loss":   cls_loss_per_sample},
                                     batch_ids=uid1, origin="train_loader")
    recon_loss_per = wl.save_signals({"train/recon_loss": recon_loss_per_sample},
                                      batch_ids=uid1, origin="train_loader")

4. Group signals for contrastive (pair-level) losses

wl.save_group_signals(
    signals={"train/contrastive_loss": contrastive_per_pair},
    group_ids=list(zip(uid1, uid2)),
    origin="train_loader",
)

wl.save_group_signals stores signals that are defined over groups of samples rather than individual ones. Each group_id is a tuple of sample UIDs. The studio can display the contrastive loss for every pair and lets you filter by pair distance to find the hardest negatives.

Tip

This example is bundled with WeightsLab (uses a synthetic dataset):

weightslab start               # 1. deploy the studio
weightslab start example --gen # 2. start the generation / anomaly demo