Metric Learning in TensorFlow & PyTorch

The same two MNIST problems — digit classification and triplet-loss metric learning — implemented twice, in TensorFlow/Keras and PyTorch, over shared data. A hands-on lab for learning PyTorch by contrast, running on an AMD GPU via ROCm.

_{Ten clean islands, one per digit — a t-SNE map of embeddings from a network that never saw a single label, trained only by comparing images in triples.}

---

The idea

Two bonus exercises from my Primeiros Passos em Machine Learning minicourse, turned into a small reusable package — and used as a personal lab to learn PyTorch starting from the Keras I already knew. Both frameworks train on the same numpy arrays and the same triplet indices, so any difference comes from the framework, not the data.

• Exercise 6 — classification (CNN). A convolutional net reads a handwritten digit and says which number it is.
• Exercise 7 — metric learning (triplet loss). A network learns a 64-number embedding per image such that same-digit images land close together — trained only from (anchor, positive, negative) triples, never from a direct label. The idea behind face recognition and similarity search.

Learning by contrast. The whole point is having the two exercises written both ways — Keras' high-level model.fit() and PyTorch's explicit training loop. My annotated study diary walks through every difference: docs/GUIA_PYTORCH.md.

Results

Measured on an AMD Radeon RX 9060 XT (RDNA4) via torch 2.9.1+rocm6.4:

Exercise	Model	Metric	Score	Train time
6 — classification	CNN (supervised)	test accuracy	98.7%	~12 s (3 epochs, GPU)
7 — metric learning	triplet embedding	kNN accuracy over embeddings	98.6%	~15 s (3 epochs, GPU)

The striking result is the second row: a network trained without ever seeing a label — only by pulling same-digit images together and pushing different ones apart — separates the digits almost as well as the fully supervised CNN.

Similarity search. Query a digit, retrieve its nearest neighbours in embedding space — all matching, by shape alone.

Confusion matrix (ex. 6). A clean diagonal; the few off-diagonal slips are the genuinely ambiguous digits.

Reproduce any of these with the scripts below — the plots land in outputs/.

Project layout

src/mnist_metric_lab/
  data.py           # self-contained MNIST loader (IDX download) + triplet sampling
  config.py         # hyperparameters, identical across backends
  devices.py        # detects the active accelerator (CUDA / ROCm / CPU)
  results.py        # shared result containers (numpy-only)
  evaluation.py     # confusion matrix, error gallery, kNN, t-SNE, similarity search
  tf_backend/       # TensorFlow/Keras: classifier.py, embedding.py
  torch_backend/    # PyTorch: models.py, classifier.py, embedding.py
scripts/
  train_classifier.py   # exercise 6, pick --backend tf|torch
  train_embedding.py    # exercise 7, pick --backend tf|torch
  benchmark.py          # runs a backend, records device + timing to outputs/
tests/                  # framework-free tests of the shared numpy core
docs/GUIA_PYTORCH.md    # the annotated study diary (Portuguese)

Only the files under torch_backend/ and tf_backend/ are framework code. Everything else — loading data, generating triplets, evaluating embeddings — is plain numpy/scikit-learn, shared by both.

Setup

Python version matters. TensorFlow and PyTorch don't publish wheels for Python 3.14 yet — use Python 3.12 (via pyenv or python3.12 -m venv). The numpy core and the tests run on 3.14; the backends don't.

cd mnist-metric-lab
python3.12 -m venv .venv && source .venv/bin/activate
pip install -e .          # installs the shared numpy/sklearn core + this package

Then add a backend:

PyTorch on AMD (RX 9060 XT / RDNA4 — the working GPU path):

# The wheel bundles its own ROCm libraries; only the in-kernel amdgpu driver is
# needed. The wheel is ~4.5 GB and unpacks under /tmp — if /tmp is a small tmpfs,
# redirect pip's temp dir to disk first (see note below).
export TMPDIR="$PWD/.pip-tmp" && mkdir -p "$TMPDIR"
pip install --index-url https://download.pytorch.org/whl/rocm6.4 torch
export HIP_VISIBLE_DEVICES=0     # pin the discrete GPU (a Ryzen iGPU shows as cuda:1)

PyTorch on NVIDIA (CUDA): pip install torch TensorFlow: pip install -r requirements-tf.txt (runs on CPU on AMD — see notes)

Usage

# Exercise 6 — classifier
python scripts/train_classifier.py --backend torch

# Exercise 7 — triplet embedding + kNN / t-SNE / similarity search
python scripts/train_embedding.py --backend torch
python scripts/train_embedding.py --backend torch --triplets 100000 --epochs 5

# Record device + timing for a backend
python scripts/benchmark.py --backends torch

Add --no-plots to skip figure generation. Run the tests (no TF/torch needed) with pytest.

GPU notes: AMD and TensorFlow

Honest expectations for this hardware, learned the hard way (full write-up in the study diary):

• PyTorch + ROCm is the working GPU path on RDNA4. The card is driven through the ordinary torch.cuda API; devices.py tells ROCm from CUDA via the HIP version. Verified here with torch 2.9.1+rocm6.4.
• pip install torch gives the CUDA build, which ignores an AMD GPU — always install from the ROCm index.
• The ~4.5 GB wheel can hit a /tmp quota (Errno 122) when /tmp is a size-capped tmpfs. Point TMPDIR at disk before installing.
• TensorFlow on AMD falls back to CPU — the official wheel is CUDA-only and tensorflow-rocm doesn't cover RDNA4. The TF backend is kept as a reference implementation; the benchmark records whichever device each framework used.

Study diary

This repo doubles as a learning artifact. The annotated walkthrough docs/GUIA_PYTORCH.md teaches PyTorch from a Keras starting point — the training loop, nn.Module, autograd, custom losses — using this project's own code, plus a field log of the AMD/ROCm setup and open questions still on my list.

Attribution

Based on bonus exercises 6 and 7 of the Primeiros Passos em Machine Learning minicourse (Semana da Física, UFES Alegre). The triplet-loss implementation follows the pattern in the official Keras documentation.

License

Licensed under the GNU Affero General Public License v3.0 (AGPL-3.0) — see LICENSE.

Copyright (C) 2026 Flávio Manoel Santos Hemerli

You may use, modify and redistribute this code, including commercially, but any derivative work — including software you run as a networked service — must be released as open source under the same AGPL-3.0 terms.