Classification, Fashion-MNIST + UMAP (PyTorch)

PyTorch classification fashion-mnist projection umap pca beta

Example: weightslab/examples/PyTorch/wl-fashion-mnist-umap/main.py

Task: 10-class Fashion-MNIST classification with a small CNN, explored with projections. The built-in parametric UMAP runs beside training; then you go back to the weights of the first 1000 steps, discard the training samples that teach the model least, lay what is left out with a PCA you write yourself, and compare the two layouts in the Projection Board and in numbers.

It is the plain-PyTorch counterpart of the projection use case, driven in stages so you can look at the board between them. Background and every option of the projection: Live 3-D Projection (Parametric UMAP).

The stages

1. Train with the live UMAP, keep a checkpoint

Nothing in the loop is projection code. Wrapping the model attaches the UMAP (it is on by default); the dict only tunes it:

model = wl.watch_or_edit(
    FashionCNN().to(device), flag="model", device=device,
    projection={"every_n_steps": 25, "graph_size": 1024},   # optional
)

train_to(1000)                                  # the UMAP fits alongside
checkpoint_step = save_checkpoint()             # weights + encoder + per-sample state

The projection hooks the input of the last Linear (fc2): the 128-D representation the classifier decides from, not the logits. The checkpoint holds the weights, the UMAP encoder and every sample’s loss at that step.

2. Reload that checkpoint

manager = ledgers.get_checkpoint_manager()
manager.load_state(manager.get_current_experiment_hash(), target_step=checkpoint_step)

Weights, optimizer, the UMAP encoder and each sample’s loss and seen-count go back to that step. Straight after stage 1 it changes nothing; it is the way back to these weights if you train on and come back to try another cut.

3. Discard what teaches little

table = sample_table()                                  # one row per training sample
easy = lowest_loss_ids(table, 0.3)                      # the 30% with the lowest loss
wl.discard_samples(easy, discarded=True)                # the loaders skip them from now on

Samples with no loss yet are never picked: an unseen sample is not an easy one. Expect the cut to come mostly from the clean classes (Trouser, Bag, the footwear) and hardly at all from the tops the model still confuses.

4. A PCA of your own

def pca_coordinates(features, dim=2):                   # plain torch, an SVD
    x = torch.as_tensor(features, dtype=torch.float32)
    x = x - x.mean(0)
    _, singular, axes = torch.linalg.svd(x, full_matrices=False)
    return (x @ axes[:dim].T).numpy()

wl.project_dataset(model, train_loader, method=pca_coordinates,
                   prefix="pca2d_step1000")

wl.project_dataset runs the loader through the model, collects the 128-D features, hands them to your function, and stores what comes back as signals//pca2d_step1000_{x,y}. The loader skips the samples you just discarded, so the PCA covers what is left. The sweep always starts from the loader’s first batch (even though training was mid-epoch on it) and does not age the model. Any algorithm returning (N, 2) or (N, 3) plugs in the same way. The name carries the step, so snapshots stay apart in the picker.

5. Compare

In the Studio, Data Exploration → Projection, open the picker and switch between umap (3-D, learned live) and pca2d_step1000 (2-D, drawn on the z = 0 plane); lasso the same region in both and compare what the grid loads.

In the script, two more views of the same two layouts. A side-by-side picture is written to <root_log_dir>/umap_vs_pca.png. And knn_purity scores each layout: for every point, the share of its k nearest neighbours in the picture that have its class (1.0 = every neighbourhood is one class, about 0.1 = ten classes thrown together). The 128-D features are scored too, as the reference both pictures summarize:

neighbours sharing a point's class (k=10), over 14000 kept samples:
  parametric UMAP (3-D) : 0.720
  your PCA (2-D)        : 0.545
  the 128-D features    : 0.858   (what both pictures summarize)

(One CPU run of the defaults; yours will differ in the last digits. UMAP keeps three dimensions and PCA two, so read the gap as a rough size, not a verdict.)

6. Retrain on what is left

The live UMAP follows the model: its encoder keeps fitting, so its layout changes with the weights. The PCA does not: it is the snapshot you took, so to see the PCA of the new model you project again under a new name (pca2d_step2000). In the same run, 1000 more steps on the remaining samples took the UMAP’s score from 0.72 to 0.84 while a fresh PCA moved from 0.55 to 0.56: the representation got more class-separable, and only the layout that learns shows it that clearly.

Every projection you compute is stored with the checkpoint of that moment (see Live 3-D Projection (Parametric UMAP)), so restoring an older checkpoint brings back the projections that existed then.

Configuration

config.yaml carries the tuning of the UMAP (projection:) and of the stages (exploration:):

projection:
  every_n_steps: 25      # fit the UMAP encoder every 25 training steps
  graph_size: 1024       # samples in the kNN graph
exploration:
  steps_initial: 1000    # train with the live UMAP, then keep a checkpoint
  discard_fraction: 0.3  # share of training samples to drop, lowest loss first
  steps_retrain: 1000    # train on what is left (0 = stop after the comparison)
  knn_k: 10              # neighbours counted by the class-purity score

Run it

weightslab start                          # terminal 1: the studio, http://localhost:8080

cd weightslab/examples/PyTorch/wl-fashion-mnist-umap
python main.py --wait --keep-serving      # terminal 2
python main.py --steps 200 100            # a quick pass: 200 initial + 100 retrain steps

--wait stops for Enter after the comparison, so you can look at the board before retraining; --keep-serving keeps the process (and the Studio connection) up at the end. The defaults take a couple of minutes on a CPU. The projection board is a beta feature.