Metadata-Version: 2.4
Name: genml_kit
Version: 0.2.0
Summary: General-purpose ML toolkit: image-classification, visual-odometry and self-supervised pre-training through one unified harness.
License: Apache-2.0
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: torch
Requires-Dist: torchvision
Requires-Dist: transformers
Requires-Dist: datasets
Requires-Dist: xgboost>=2.0
Requires-Dist: scikit-learn
Requires-Dist: pillow
Requires-Dist: tensorboard
Provides-Extra: gcs
Requires-Dist: google-cloud-storage; extra == "gcs"
Provides-Extra: s3
Requires-Dist: boto3>=1.28; extra == "s3"
Provides-Extra: lora
Requires-Dist: peft>=0.7.0; extra == "lora"
Provides-Extra: timm
Requires-Dist: timm>=1.0; extra == "timm"
Provides-Extra: uvito
Requires-Dist: segmentation_models_pytorch>=0.3.0; extra == "uvito"
Provides-Extra: rl
Requires-Dist: gymnasium>=0.29; extra == "rl"
Requires-Dist: pygame>=2.1; extra == "rl"
Requires-Dist: imageio[ffmpeg]>=2.31; extra == "rl"
Provides-Extra: all
Requires-Dist: genml_kit[gcs,lora,rl,s3,timm,uvito]; extra == "all"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: yapf; extra == "dev"
Dynamic: license-file

# genml_kit

[![CI](https://github.com/davidel/genml_kit/actions/workflows/ci.yml/badge.svg)](https://github.com/davidel/genml_kit/actions/workflows/ci.yml)

A general-purpose toolkit for image classification. Supports self-supervised
pre-training, supervised fine-tuning, and XGBoost ensemble inference — all
from the command line.

`genml_kit` is domain-agnostic: point it at any HuggingFace dataset, local
`ImageFolder` tree, or timm/HuggingFace backbone. It grew out of a
skin-lesion classification project; the dermoscopy-specific workflow (dataset
preparation, tuned recipes) is documented separately in
[`scdiag/README.md`](scdiag/README.md). The drone visual-odometry front-end
(similarity estimation between oblique aerial frames) is documented separately
in [`vo/README.md`](vo/README.md).

## Contents

- [Why genml_kit?](#why-genml_kit)
- [How It Works](#how-it-works)
  - [The Method Lifecycle](#the-method-lifecycle)
- [Installation](#installation)
- [Quick Start](#quick-start)
- [Pre-Training Guide](#pre-training-guide)
- [Fine-Tuning Guide](#fine-tuning-guide)
- [Inference Guide](#inference-guide)
- [Tips & Pitfalls](#tips--pitfalls)
- [Gradient Monitor](#gradient-monitor)
- [Custom Losses](#custom-losses)
- [Custom Models](#custom-models)
- [References And Further Reading](#references-and-further-reading)
- [Development](#development)
- [License](#license)

## Why genml_kit?

Practitioners face two recurring problems: **labeled data is scarce** and
**off-the-shelf models are not domain-specific**. A ViT pre-trained on
ImageNet can classify cats and dogs, but specialized imagery — medical,
satellite, industrial, scientific — looks nothing like natural photos, and
the feature distributions are fundamentally different.

genml_kit solves this with a two-stage pipeline:

1. **Pre-train** on large, often unlabeled image collections (or a labeled
   superset) to learn domain-appropriate visual features.
2. **Fine-tune** on your smaller labeled dataset, starting from those
   pre-trained features instead of random initialization.

This consistently outperforms training from scratch, especially when your
labeled dataset has fewer than ~5 000 images. The tool also supports ensemble
inference with XGBoost on top of the learned features, which can squeeze out
additional performance for deployment.

> **Concrete example:** the dermoscopy workflow this toolkit was extracted
> from — preparing HAM10000 / ISIC / Derm1M corpora, per-method
> hyperparameters, and a worked pre-train → fine-tune run — is documented in
> [`scdiag/README.md`](scdiag/README.md).

## How It Works

```
┌─────────────────────────────────────────────────────────────────┐
│                        Pre-Training                             │
│  Unlabeled/labeled images                                       │
│  ──────────────────────────────────►  Encoder with learned      │
│  SimMIM / I-JEPA / SupCon              visual features          │
└────────────────────────────┬────────────────────────────────────┘
                             │  encoder weights
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│                       Fine-Tuning                               │
│  Small labeled dataset  ──────────►  Trained classifier         │
│  + pre-trained encoder                for your task             │
└────────────────────────────┬────────────────────────────────────┘
                             │  backbone features
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│                     (Optional) Ensemble                         │
│  Backbone features  ──────────►  XGBoost on top of the         │
│                                   learned representations      │
└─────────────────────────────────────────────────────────────────┘
```

**Pre-training** teaches the model to understand the target imagery — textures,
boundaries, colour patterns, and spatial relationships. **Fine-tuning** adapts
that understanding to your specific classification task (e.g. disease vs.
healthy, defective vs. passing). **Ensemble inference** (optional) trains a
tree-based model on the same features, which sometimes generalises better
than a linear head for small datasets.

### A Useful Mental Model

The encoder turns an image into a vector of features. During pre-training, we
choose an artificial task whose answer can be obtained from the images (or,
for SupCon, from their labels). The encoder learns parameters $`\theta`$ that
make this task easy. During fine-tuning, a classifier is attached to the
encoder and the whole model, or a selected part of it, is adapted to the real
labels:

```text
image x  ──► encoder f_theta(x)  ──► classifier g_phi  ──► class probabilities
                    │                         │
             reusable features          task-specific boundary
```

Here $`x`$ is the input image, $`f_\theta(x)`$ the feature vector computed by
the encoder with parameters $`\theta`$, and $`g_\phi`$ the classifier head with
parameters $`\phi`$.

Pre-training and fine-tuning are not two names for the same job. Pre-training
shapes a useful representation; fine-tuning decides how that representation
should be used for the target labels. If the pre-training data is visually
related to the target data, this gives the classifier a much better starting
point than random initialization. If the domains are very different, use a
smaller learning rate for the backbone and validate carefully.

The pre-training methods are introduced above with links to their papers (see
the [Pre-Training Guide](#pre-training-guide)). The
[timm documentation](https://huggingface.co/docs/timm/index) and the
[Hugging Face image classification guide](https://huggingface.co/docs/transformers/tasks/image_classification)
are useful references when selecting a backbone or processor.

### The Method Lifecycle

`genml-kit-train` is a *branchless* driver: it never inspects which
`--method` or `--pipeline` it is running.  Instead, every objective registers
itself through the `Method` registry and declares the same four lifecycle
hooks, which the driver calls in one fixed order:

| # | Hook | Purpose |
|---|---|---|
| 1 | `prepare_transforms(args, device)` | Resolve any transforms the method needs (e.g. the HF processor's normalization) **before** the loaders are built. |
| 2 | `wire_data(args, pipeline)` | Consume pipeline-built data attributes: label space, class weights, criterion. Runs **before** the model so the HF head can be sized from the label space. |
| 3 | `build_model(args, device)` | Construct the model. Every method applies LoRA (`--lora`), source-checkpoint weights (`--source_checkpoint`), freeze patterns (`--freeze`) and gradient checkpointing (`--grad_checkpoint`) through the same `_apply_model_extras` step — so all four flags work for **every** objective, not just classification. |
| 4 | `post_train(args, pipeline, device, result)` | Optional post-training stage, called after `BaseTrainer.run()` returns (skipped when the run was interrupted). |

Two ordering invariants are enforced by the hook order itself, not by
convention: transforms resolve before loaders consume them, and the label
space exists before the classification head is constructed.

Adding a new objective means implementing the hooks you need (only
`build_model` and `train_step` are abstract); the driver never changes.
`ClassificationMethod` additionally implements `post_train` to run the
optional XGBoost shallow-head stage; any future classifier-style method can
reuse the same hook.

## Package Layout

| Module | Contents |
|---|---|
| `genml_kit.training` | `train.py` and `infer.py` CLI harnesses, `optim_factory.py` (optimizers, LLRD, schedulers), `model_utils.py` (loading, freezing, feature extraction), `param_align.py`, `eval.py` + `metrics.py`, `grad_monitor.py`, `train_reporting.py`, `tta.py`, `xgb_utils.py` + `xgb_pipeline.py`, `classifiers/` (pluggable heads) |
| `genml_kit.methods` | objective-side registry (classification, vo_pair, SimMIM, I-JEPA, DINO, BYOL, SupCon) |
| `genml_kit.losses` | shared loss library (focal, NT-Xent contrastive, DINO self-distillation, BYOL) |
| `genml_kit.augmentations` | objective-level transforms (dual-view, multi-crop) |
| `genml_kit.models` | model/processor registry; `timm/`, `convvit/`, `uvito/`, `cls_model_wrapper/` backends; `processors/base.py` |
| `genml_kit.datasets` | `hf_proxy.py` (HuggingFace → PyTorch bridge), `image_folder.py`, `ensemble.py`, `field_dataset.py`, `balanced_sampler.py`, `weighted_sampler.py`, `retry.py`, `vo_pairs.py` (synthetic VO pairs) |
| `genml_kit.geometry` | `similarity.py` (2-D similarity algebra, batched differentiable Umeyama, backward-map warp) — the math behind the visual-odometry front-end lives in [`vo/README.md`](vo/README.md) |
| `genml_kit.io` | `checkpointing.py` (atomic saves, LoRA state, remote fetch), `storage_utils.py` (S3 / GCS / R2) |
| `genml_kit.utils` | glog-style logging, CLI arg groups, seeding, signal handling, GPU info, tables, external `.py` script loading, `image_dump`, transformer init helpers |

## Installation

```bash
pip install genml_kit

# With timm model support:
pip install "genml_kit[timm]"

# With GCS checkpoint sync:
pip install "genml_kit[gcs]"

# With AWS S3 / Cloudflare R2 checkpoint sync (both use boto3):
pip install "genml_kit[s3]"

# With LoRA fine-tuning:
pip install "genml_kit[lora]"

# With UVito model support:
pip install "genml_kit[uvito]"

# Everything above in one shot (gcs, s3, lora, timm, uvito):
pip install "genml_kit[all]"
```

**Requirements:** Python ≥ 3.9, PyTorch, torchvision, transformers, datasets,
NumPy, scikit-learn ≥ 1.3, XGBoost ≥ 2.0, Pillow, tensorboard.

## Quick Start

The fastest way to get started:

```bash
# Fine-tune a ViT on an image-classification dataset (5 epochs, ~2 minutes on GPU)
genml-kit-train --model google/vit-base-patch16-224 \
                --dataset cifar10 \
                --label_column label \
                --epochs 5 \
                --batch_size 32 \
                --lr 3e-5 \
                --image_size 224
```

HuggingFace ViT backbones use fixed 224x224 position embeddings and do not
interpolate them, so they require `--image_size 224`.  The toolkit default of
448 targets size-flexible backbones such as ConvNeXt, timm EVA, or ConvViT,
which pool or interpolate the token sequence to the input size.

This downloads the model and dataset from HuggingFace, trains for 5 epochs,
and saves `genml_kit_latest.pt` and `genml_kit_best.pt`. See
[Pre-Training Guide](#pre-training-guide) below for the full pipeline
(starting with pre-training before fine-tuning).

---

## Pre-Training Guide

Pre-training learns general visual features from large datasets *before* you
fine-tune on your specific task. This is especially valuable wherever labeled
data is expensive to obtain but raw images are available in bulk — medical
imaging, remote sensing, industrial inspection, scientific imaging.

genml_kit supports five pre-training methods, each with different strengths:

### Choosing Your Method

| Method | Needs labels? | Best when… | Key idea |
|---|---|---|---|
| **SimMIM** | No | You have large unlabeled datasets; want a simple, proven approach | Mask 60% of image patches, train the model to reconstruct the raw pixels |
| **I-JEPA** | No | You want faster training and better downstream transfer than SimMIM | Predict *representations* of masked regions, not raw pixels — avoids learning noise |
| **SupCon** | Yes | You have labels and want representations that cluster by class | Pull same-class images together, push different classes apart in feature space |
| **DINO** | No | You want strong features from a modest amount of data; proven, robust default | Self-distillation: a student matches the output distribution of its own EMA copy (the teacher) |
| **BYOL** | No | Similar goals to DINO with a simpler pipeline (two views, no multi-crop) | Online network predicts the output of its own EMA target network |

### How Each Method Works

The methods differ mainly in what they call a correct answer. SimMIM asks for
pixels, I-JEPA asks for features, DINO and BYOL ask a network to predict the
output of its own slowly-moving copy, and SupCon asks for relative positions in
feature space. That distinction matters: pixel reconstruction can spend effort
reproducing colour and high-frequency detail, while the other methods spend
effort building invariant, separable representations.

**SimMIM** (Masked Image Modelling): Randomly masks ~60% of image patches and
trains a lightweight decoder to reconstruct the original pixels. The encoder
must learn to understand textures, boundaries, and spatial context from just
40% of the image. Think of it as a "fill in the blanks" exercise for vision
models. Good default choice when you have lots of unlabeled images.
Reference: Xie et al., [SimMIM: A Simple Framework for Masked Image Modeling](https://arxiv.org/abs/2111.09886), CVPR 2022.

For an image split into patches $`x_1, \dots, x_N`$, let $`\mathcal{M}`$ be
the set of masked patch indices and $`\hat{x}_i`$ the decoder prediction.
SimMIM minimizes the mean squared reconstruction error over masked patches:

$$
\large
\mathcal{L}_{\text{MIM}} = \frac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \left\| \hat{x}_i - x_i \right\|_2^2
$$

where:

- $`x_i`$ — the pixels of the $`i`$-th original patch;
- $`\hat{x}_i`$ — the decoder's reconstruction of that patch;
- $`\mathcal{M}`$ — the set of masked patch indices, and $`|\mathcal{M}|`$ its
  size (how many patches are masked);
- $`\lVert \cdot \rVert_2^2`$ — the squared Euclidean ($`L_2`$) norm, i.e. the
  sum of squared pixel differences over the patch.

Only masked patches contribute to the loss; otherwise copying visible pixels
would make the task too easy. A higher
`--mask_ratio` supplies less context and creates a harder task, but an
excessively high ratio can make reconstruction ambiguous.

**I-JEPA** (Joint-Embedding Predictive Architecture): Also masks patches, but
instead of reconstructing pixels, it predicts the *latent representation* of
the masked region from the visible context. This avoids wasting capacity on
pixel-level noise (e.g. exact JPEG compression artifacts) and learns more
transferable features. Uses a teacher–student setup with EMA momentum ramping.
Reference: Assran et al., [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://arxiv.org/abs/2301.08243), CVPR 2023.

Let $`f_\theta`$ be the student encoder and $`f_\xi`$ the teacher encoder;
$`\theta`$ and $`\xi`$ are their respective parameter vectors (the subscripts
name the network a symbol belongs to). The predictor $`q_\theta`$ receives the
visible context and predicts a target representation $`z_j = f_\xi(x_j)`$ for a
masked region. The objective is regression in representation space:

$$
\large
\mathcal{L}_{\text{IJEPA}} = \frac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \left\| q_\theta(f_\theta(\text{context}))_i - \mathop{\mathrm{stopgrad}}\left(f_\xi(x_i)\right) \right\|_2^2
$$

where:

- $`x_i`$ — the $`i`$-th masked patch (the prediction target);
- $`f_\theta(\text{context})`$ — the student's encoding of the *visible* patches;
- $`q_\theta(\cdot)_i`$ — the predictor's output for the $`i`$-th masked location;
- $`f_\xi(x_i)`$ — the teacher's representation of the masked patch;
- $`\mathrm{stopgrad}(\cdot)`$ — treat the argument as a constant: gradients do
  not flow through it during the student update;
- $`\lVert \cdot \rVert_2^2`$ — squared Euclidean distance in representation space;
- $`\mathcal{M}`$ — the set of masked patch indices.

The teacher is not optimized by backpropagation; it follows the student with
an exponential moving average:

$$
\large
\xi \leftarrow m \, \xi + (1 - m) \, \theta
$$

where $`m`$ is the momentum (`--teacher_momentum`) and the arrow denotes
assignment: the new teacher parameters are a convex blend of the old teacher
and the current student. A high $`m`$ changes the teacher slowly and gives
more stable targets. The asymmetric teacher update and masking are important:
two networks simply trained to copy each other could collapse to a constant
vector (the teacher-follows-student trick and $`\mathrm{stopgrad}`$ come from
Grill et al., [BYOL](https://arxiv.org/abs/2006.07733), NeurIPS 2020).

Three mechanism details complete the picture:

- **Block masking.** Target regions are not scattered single patches: the mask
  is composed of a few contiguous rectangular blocks (4 blocks of 6×6 patches
  in the current implementation). Contiguous holes are much harder to fill in
  from context than isolated ones, so the task rewards spatial understanding
  rather than high-frequency texture interpolation.
- **Two crop scales.** The *context* (source) view is a small random crop (50%
  of the image side) and the *target* view is a larger one (85%). The predictor
  must infer representations of a region it sees only through a narrow window,
  which encourages semantic rather than layout-memorized features.
- **Momentum ramp.** $`m`$ moves from `--teacher_momentum` (0.996) toward
  `--teacher_final_momentum` (1.0) as training progresses: as the encoder
  matures, the teacher freezes progressively, stabilizing the targets late in
  training.

**SupCon** (Supervised Contrastive Learning): Uses labels to define "positive"
pairs (same class) and "negative" pairs (different classes). The loss pulls
features of same-class images together and pushes different-class features
apart. Produces a feature space where similar images naturally cluster.
Requires a `ContrastiveEncoder` (backbone + projection head) and balanced
batch sampling to ensure each batch has enough same-class pairs.
Reference: Khosla et al., [Supervised Contrastive Learning](https://arxiv.org/abs/2004.11362), NeurIPS 2020.

An image $`x_i`$ is mapped by the encoder-and-projection-head
$`f_\theta`$ to a vector $`f_\theta(x_i)`$, which is rescaled to unit length
(normalized) so that comparisons are scale-invariant:

$$
\large
z_i = \frac{f_\theta(x_i)}{\lVert f_\theta(x_i) \rVert_2}
$$

where $`\lVert v \rVert_2 = \sqrt{\sum_d v_d^2}`$ is the Euclidean ($`L_2`$)
length of a vector $`v`$, so every $`z_i`$ lies on the unit hypersphere.
The similarity of two examples $`i`$ and $`j`$ is then the dot product of
their normalized projections, divided by the temperature $`\tau`$
(`--temperature`):

$$
\large
s_{ij} = \frac{z_i^\top z_j}{\tau}
$$

where $`z_i^\top z_j`$ is the dot product (cosine of the angle between the two
unit vectors) and $`\tau`$ rescales the similarities before the softmax.

Fix one example $`i`$, called the *anchor*. Its *positive set* is

$$
\large
\mathcal{P}(i) = \{\, j : j \neq i \text{ and } y_j = y_i \,\}
$$

i.e. all other examples in the batch that share the anchor's class label
$`y_i`$; every remaining example is a *negative*. SupCon averages the
log-softmax probability that the anchor assigns to each of its positives:

$$
\large
\mathcal{L}_i = -\frac{1}{|\mathcal{P}(i)|} \sum_{p \in \mathcal{P}(i)} \log \left( \frac{\exp(s_{ip})}{\sum_{a \neq i} \exp(s_{ia})} \right)
$$

$$
\large
\mathcal{L} = \frac{1}{B} \sum_{i=1}^{B} \mathcal{L}_i
$$

where:

- $`s_{ip}`$ — similarity between the anchor and positive $`p`$ (above);
- $`\sum_{a \neq i}`$ — sum over all other examples in the batch, positives
  *and* negatives (the softmax denominator);
- $`|\mathcal{P}(i)|`$ — the number of positives for the anchor;
- $`B`$ — the batch size; the total loss averages the per-anchor terms.

Minimizing $`\mathcal{L}`$ pushes $`s_{ip}`$ up for same-class pairs and down
for different-class pairs. A lower temperature makes the distribution sharper:
this
can help separate hard negatives but can also make optimization less stable.
`--samples_per_class` matters because an anchor needs another example of its
class to have a positive. A class represented once contributes no useful
SupCon term for that anchor.

#### Temperature

The division by $`\tau`$ (`--temperature`, default 0.07) above, and the two
temperatures used by DINO below, all do the same job. Softmax over logits
$`s`$ becomes

$$
\large
\mathrm{softmax}(s)_j = \frac{\exp(s_j / \tau)}{\sum_k \exp(s_k / \tau)}
$$

A small $`\tau`$ divides every logit by a tiny number, so exponentials grow
very fast and the distribution concentrates on the maximum (sharp); a large
$`\tau`$ flattens it toward uniform. Sharpening amplifies small differences in
similarity, which speeds learning, but too-small $`\tau`$ makes gradients
vanish for non-maximum classes and destabilizes training. SupCon uses one
moderate value; DINO uses a small teacher temperature (sharper *targets*,
they are meant to be confident) and a larger student temperature (softer
*predictions*, easier to move).

**DINO** (self-**di**stillation with **no** labels): Caron et al.,
[Emerging Properties in Self-Supervised Vision Transformers](https://arxiv.org/abs/2104.14294), ICCV 2021.
One backbone is instantiated twice. The **student** $`f_\theta`$ is trained by
gradient descent; the **teacher** $`f_\xi`$ is the EMA copy of the student
($`\xi \leftarrow m\,\xi + (1-m)\,\theta`$, as in I-JEPA above, ramped from
`--dino_momentum` 0.996 to `--dino_final_momentum` 1.0). Each image is cut
into 2 **global** crops plus `--dino_local_num` (8) small **local** crops:
the teacher sees only global views, the student sees all of them, so the
student must produce globally consistent features from partial views.

Both towers end in an MLP projection head whose output is turned into a
probability vector over the projection dimension. The teacher's distribution
is *sharpened* by its small temperature and *centered* by a running mean
$`c \in \mathbb{R}^{D}`$ updated once per step,

$$
\large
c \leftarrow m_c \, c + (1 - m_c) \cdot \frac{1}{B}\sum_{i=1}^{B} f_\xi(x_i)
$$

(`--dino_center_momentum` 0.9), and the loss is the cross-entropy between the
teacher's processed distribution $`p_\xi`$ and the student's $`p_\theta`$,
averaged over all (teacher global, student crop) pairings:

$$
\large
\mathcal{L} = - \sum_{j} p_{\xi,j} \, \log p_{\theta,j}
$$

Centering and sharpening are the two anti-collapse mechanisms: centering
cancels any component the teacher produces for *every* input (a constant
output carries no information and would otherwise still lower the loss), and
sharpening makes the teacher confident, forcing the student to truly reproduce
its argmax rather than drifting toward a flat, uninformative distribution.
Because the targets come from the teacher and the teacher follows the student,
this is self-distillation — there are no external labels anywhere.

**BYOL** (**B**ootstrap **Y**our **O**wn **L**atent): Grill et al.,
[Bootstrap Your Own Latent](https://arxiv.org/abs/2006.07733), NeurIPS 2020.
The same two-tower idea as DINO but with two augmented views and no
multi-crop. The **online** network (encoder + projector + predictor MLP) and
the **target** network (encoder + projector, EMA copy) process the two views;
only the online tower receives gradients. With $`z_1, z_2`$ the target
projections of the two views and $`p_1, p_2`$ the online *predictions* of
those projections, the loss is the symmetric cosine regression

$$
\large
\mathcal{L} = 2 - \cos(p_1, z_2) - \cos(p_2, z_1)
$$

which is 0 when predictions perfectly match target projections and reaches 4
when they point in opposite directions; $`\cos(u,v) = u^\top v / (\lVert u
\rVert_2 \lVert v \rVert_2)`$ is the cosine of the angle between the two
vectors. Two asymmetries keep this from collapsing to a constant: the
predictor exists only on the online side (predicting a target is easier than
being one), and the targets move slowly under the EMA — exactly the
teacher-follows-student trick noted under I-JEPA.

### Typical Hyperparameters

These are reasonable starting points. Tune from here based on your dataset
size and GPU memory:

| Parameter | SimMIM | I-JEPA | DINO | BYOL | SupCon |
|---|---|---|---|---|---|
| `--image_size` | 448 | 448 | 224 | 224 | 448 |
| `--batch_size` | 32 | 32 | 64 | 64 | 64 |
| `--lr` | 1e-4 | 1e-4 | 5e-4 | 5e-4 | 1e-4 |
| `--epochs` | 200 | 200 | 100 | 100 | 100 |
| `--scheduler` | CosineAnnealingLR | CosineAnnealingLR | CosineAnnealingLR | CosineAnnealingLR | CosineAnnealingLR |
| `--amp_dtype` | bfloat16 | bfloat16 | bfloat16 | bfloat16 | bfloat16 |
| `--mask_ratio` | 0.6 | — | — | — | — |
| `--teacher_momentum` | — | 0.996→1.0 | — | — | — |
| `--dino_momentum` | — | — | 0.996→1.0 | — | — |
| `--byol_momentum` | — | — | — | 0.996→1.0 | — |
| `--temperature` | — | — | — | — | 0.07 |
| `--samples_per_class` | — | — | — | — | 16 |

**Tips:**
- Start with 200 epochs for SimMIM/I-JEPA. DINO, BYOL, and SupCon converge
  faster (~100).
- `--temperature 0.07` is the standard from the original SupCon paper.
  Lower = sharper contrastive distribution; try 0.05–0.1.
- DINO and BYOL train on smaller crops (224) with a larger LR because their
  loss acts on a small projection space rather than every pixel; the defaults
  above match the method-specific `--dino_*` / `--byol_*` flags.
- `--samples_per_class 16` with `--batch_size 64` gives 4 classes per batch
  on a 7-class dataset. batch_size need not be divisible by
  samples_per_class; the remainder is spread evenly across the per-class
  groups so each batch still contains exactly batch_size samples.
- Use `--amp_dtype bfloat16` if your GPU supports it (Ampere+). Otherwise
  `float16` with GradScaler works too.

### Example: Full Pre-Training Pipeline

```bash
# Step 1: Pre-train with SimMIM on two large image collections.
# v4.2: pre-training is a *configuration* of the unified genml-kit-train
# CLI (--pipeline images + --method simmim), not a separate binary.
genml-kit-train --pipeline images --method simmim \
                --model convvit \
                --datasets "imagefolder/raw-photos" "imagefolder/more-photos" \
                --image_size 448 \
                --batch_size 32 \
                --epochs 200 \
                --lr 1e-4 \
                --scheduler CosineAnnealingLR \
                --sched_arg T_max=200 --sched_arg eta_min=1e-6 \
                --amp_dtype bfloat16 \
                --checkpoint ./checkpoints/convvit_simmim

# Step 2: Fine-tune on your labeled dataset
genml-kit-train --model convvit \
                --dataset my-org/labeled-photos \
                --label_column category \
                --source_checkpoint ./checkpoints/convvit_simmim_latest.pt \
                --epochs 100 \
                --lr 3e-5 \
                --batch_size 32 \
                --amp_dtype bfloat16
```

### Example: Supervised Contrastive Pre-Training

```bash
genml-kit-train --pipeline images --method supcon \
                --model convvit \
                --datasets my-org/labeled-photos \
                --label_column category \
                --image_size 448 \
                --batch_size 64 \
                --samples_per_class 16 \
                --proj_dim 128 \
                --temperature 0.07 \
                --epochs 100 \
                --lr 1e-4 \
                --amp_dtype bfloat16 \
                --checkpoint ./checkpoints/convvit_supcon

# Then fine-tune as above with --source_checkpoint ./checkpoints/convvit_supcon_latest.pt
```

### Pre-Training CLI Reference

| Argument | Default | Description |
|---|---|---|
| `--method` | `classification` | Objective. The pre-training examples below pass `--method simmim` explicitly. Choices: `classification`, `simmim`, `ijepa`, `dino`, `byol`, `supcon`, `vo_pair`. |
| `--model` | `convvit` | Model name registered in genml_kit or HuggingFace model ID. |
| `--datasets` | (required) | Space-separated dataset names or local paths. |
| `--cache_dir` | `None` | HuggingFace cache directory for downloads. |
| `--remote_checkpoint` | `None` | Remote URI for checkpoint sync (`gs://BUCKET/PREFIX`, `r2://BUCKET/PREFIX`, or `s3://BUCKET/PREFIX`). |
| `--hf_token` | `None` | HuggingFace token for gated datasets (or set `HF_TOKEN` env var). |
| `--image_column` | auto-detected | Explicit HF image column name. |
| `--label_column` | `None` | Explicit HF label column name. Required by `--method supcon` if non-standard. |
| `STRICT_DATASETS` (env) | `False` | Abort on first dataset-loading failure instead of skipping. |
| `--image_size` | `448` | Input image size (square). |
| `--batch_size` | `32` | Per-GPU batch size. |
| `--seed` | `None` | Explicit RNG seed: pins the shuffling and enables full determinism (cuDNN deterministic kernels, benchmark off). Omit (default) to keep runs fast — RNG streams are still seeded from `$GENML_KIT_SEED` (default 42). See [Reproducibility](#reproducibility). |
| `--epochs` | `5` | Number of epochs. The recipes below override this for pre-training. |
| `--lr` | `3e-5` | Peak learning rate (base). The recipes below override this. |
| `--amp_dtype` | `None` | Mixed precision: `float16` or `bfloat16`. Omit to disable. |
| `--num_workers` | `4` | DataLoader worker processes. |
| `--device` | auto-detect | Device: `cpu`, `cuda`, or `cuda:INDEX`. |
| `--resume` | `True` | Auto-resume from latest checkpoint. Use `--no-resume` to disable. |
| `--state_save` | `opt,sched,amp` | States to save: `opt`, `sched`, `amp`, `none`. |
| `--state_load` | `opt,sched,amp` | States to restore on resume: `opt`, `sched`, `amp`, `none`. |
| `--checkpoint` | (required) | Checkpoint path prefix (saves `_latest.pt` and `_best.pt`). |
| `--log_level` | `INFO` | Minimum logging level. |
| `--log_targets` | `STDERR` | Comma-separated log destinations. `STDERR` logs to standard error; any other entry is a log file path (appended). Example: `STDERR,/tmp/train.log` logs to both. |
| `--grad_monitor` | `-1` | Log gradient statistics every N steps; `-1` disables. See [Gradient Monitor](#gradient-monitor). |
| `--norm_history` | `0` | Keep last N norm snapshots per parameter for trend analysis. |
| `--trend_top_n` | `10` | Show top N params in trend table by abs change %. `0` = show all. |
| `--grad_clip` | `1.0` | Maximum gradient norm for clipping. `0` disables. |
| `--lr_group` | `None` | Per-parameter-group learning rates (repeatable). Format: `"REGEX=LR"`. |
| `--llrd_decay` | `None` | Layer-wise learning rate decay factor. |
| `--vis_every` | `0` | Log reconstruction visualisation every N steps (SimMIM only). |
| `--save_every` | `500` | Save checkpoint every N optimizer steps. 0 disables. |
| `--model_arg` | `{}` | Override model configuration (repeatable). For `--model cls_model_wrapper:<hf_name>` the classifier head travels here too: `--model_arg classifier=mlp:hidden=512,dropout=0.3`. |
| `--proc_arg` | `{}` | Override processor configuration (repeatable). |
| `--optimizer` | `AdamW` | `torch.optim` optimizer class name or `.py` script path. |
| `--opt_arg` | `{}` | Extra optimizer kwargs (repeatable). |
| `--scheduler` | `None` | `torch.optim.lr_scheduler` class name or `.py` script path. |
| `--sched_arg` | `{}` | Extra scheduler kwargs (repeatable). |
| `--source_checkpoint` | `None` | Path to source checkpoint to absorb parameters from. |
| `--param_rename` | `None` | Regex-based key rename patterns (`SEARCH;REPLACE`). |
| `--grad_checkpoint` | `False` | Enable gradient checkpointing. Reduces activation memory by ~40-50% at the cost of ~25-35% more compute per step. Enables larger batch sizes. See [Gradient Checkpointing](#gradient-checkpointing). |

**Defaults shown are the generic `genml-kit-train` defaults; the
pre-training recipes below override `--method`, `--epochs`, `--lr` for their
specific objective.** The pre-training examples invoke
`genml-kit-train --pipeline images --method simmim ...` explicitly.

**SupCon-specific arguments:**

| Argument | Default | Description |
|---|---|---|
| `--proj_dim` | `128` | Output dimensionality of the projection head. |
| `--proj_hidden` | `None` | Hidden layer size of the projection MLP. `None` = single linear layer. |
| `--temperature` | `0.07` | NT-Xent temperature. Lower = sharper contrastive distribution. |
| `--samples_per_class` | `16` | Samples per class in each batch. Batch size need not be divisible by this; the remainder is spread evenly across the per-class groups. |

### Dataset Ensemble

With `--pipeline images`, `genml-kit-train` stitches multiple datasets into
a single pre-training corpus (via `--datasets`). This is useful because no
single dataset is large enough for effective pre-training on its own.

Supported dataset types:
- **HuggingFace datasets** — any HF dataset ID that returns decoded image
  data (e.g. `cifar10`, `food101`). Gated datasets require `--hf_token` or
  `HF_TOKEN`.
- **Local image directories** — pass a path to a folder of images
  (ImageFolder format).

Datasets are loaded lazily (only when first accessed). By default, datasets
that fail to load are logged and skipped (best-effort mode). Set the
`STRICT_DATASETS` environment variable to a truthy value (`1`, `true`,
`yes`, `on`) to abort on the first failure instead.

Images that cannot be decoded are skipped with a warning — this prevents a
single corrupted file from blocking an entire pre-training run.

#### Label Validation

When using a label-aware method (`--method supcon`), the ensemble validates
that every dataset supports labels. Datasets without a label column cause a
clear error *before* training begins, not a cryptic runtime failure mid-epoch.

Labels are automatically remapped to a shared global label space across all
datasets, so mixing datasets with overlapping but differently-named classes
works transparently.

### Preparing Datasets

Some datasets store images inside zip archives or need custom preprocessing
before they can be used for pre-training. The toolkit consumes any local
ImageFolder directory, so a small preparation script is all it takes:

```bash
python my_prepare_script.py --output_dir ./prepared_images
```

Then use the extracted directory as a local dataset:

```bash
genml-kit-train --pipeline images --method simmim \
                --datasets ./prepared_images ./other-images \
                --image_size 448 --batch_size 32 ...
```

For a concrete worked example of such a preparation script, see
`scdiag/scripts/prepare_derm1m.py` and `scdiag/scripts/prepare_ham10000.py`
in the repository (dermoscopy corpora).

---

## Fine-Tuning Guide

After pre-training (or directly, if you skip pre-training), fine-tune a
classifier on your labeled dataset.

### What Fine-Tuning Is Changing

Suppose the encoder produces the feature vector $`h = f_\theta(x)`$. A linear
classification head computes logits $`a = W h + b`$, and softmax turns them
into probabilities:

$$
\large
p(y=c \mid x) = \frac{\exp(a_c)}{\sum_k \exp(a_k)}
$$

where:

- $`h \in \mathbb{R}^{D}`$ — the encoder's feature vector for image $`x`$;
- $`W \in \mathbb{R}^{C \times D}`$, $`b \in \mathbb{R}^{C}`$ — the head's weight
  matrix and bias ($`C`$ = number of classes, $`D`$ = feature dimension);
- $`a_c`$ — the logit (raw score) for class $`c`$;
- $`p(y=c \mid x)`$ — the model's probability that the class of $`x`$ is
  $`c`$; the softmax exponentials make the scores positive and summing to one.

Training minimizes cross-entropy between the model's predicted distribution
$`p`$ and the target distribution $`q`$ over the $`C`$ classes,

$$
\large
\mathcal{L}_{\text{CE}}(p, q) = -\sum_{c=1}^{C} q_c \log p_c
$$

For a hard label the target $`q`$ is a one-hot vector, so only the term of the
true class survives and the loss reduces to the familiar $`-\log p(y \mid x)`$.
genml_kit generalizes this in three directions via `CombinedFocalLoss`:

- **Soft targets.** With Mixup (below) $`q`$ is the blend
  $`\lambda\, y_i + (1-\lambda)\, y_j`$, and the sum keeps both classes'
  terms weighted by $`\lambda`$ and $`1-\lambda`$.
- **Label smoothing** (via `--loss_args '{"label_smoothing": ...}'`):
  the one-hot target is replaced by
  $`q_c = 1 - \varepsilon + \varepsilon/C`$ for the true class and
  $`\varepsilon/C`$ for the others. The model is no longer asked to drive
  $`p_c \to 1`$ exactly, which reduces overconfidence and improves
  calibration on validation data.
- **Class weights** (`--class_multipliers`): each term of the sum is scaled by
  a per-class factor $`w_c`$, so rare classes contribute larger gradients.

The new classifier head is normally initialized from scratch because
its output
size depends on the target classes. When loading a pre-training checkpoint,
the useful part to transfer is the encoder; an unused SupCon projection head
should not be copied into the classifier.

The practical choice is how much of $`\theta`$ to update:

- **Full fine-tuning** updates encoder and head. It gives the model the most
  freedom, but needs enough data and a conservative learning rate.
- **Frozen-backbone training** updates only the head. It is a useful baseline
  for small datasets and shows how much information the representation holds.
- **LLRD** updates all layers but gives early layers smaller learning rates.
  This is often a good compromise when the pre-training domain differs from
  the fine-tuning domain (e.g. ImageNet to specialized imagery).
- **LoRA** freezes the original matrices and learns small low-rank updates.
  It is useful when GPU memory or labeled data is limited.

Compare these strategies on the same validation split. The lowest training
loss is not necessarily the best model: monitor macro-F1, balanced accuracy,
weighted F1, and per-class precision and recall, especially for minority
classes.

### Evaluation Metrics

Validation reports include:

- **Top-1 accuracy:** the percentage of validation images assigned the correct
  class by the highest-probability prediction.
- **Precision:** for a class, the fraction of images predicted as that class
  that truly belong to it:
  $`\text{precision} = TP / (TP + FP)`$.
- **Recall:** for a class, the fraction of images belonging to that class that
  are predicted correctly: $`\text{recall} = TP / (TP + FN)`$.
  Here $`TP`$ (true positives) counts images of the class that were correctly
  predicted as that class, $`FP`$ (false positives) counts images of *other*
  classes wrongly predicted as this class, and $`FN`$ (false negatives) counts
  images of this class wrongly predicted as something else.
- **F1:** the harmonic mean of precision and recall:
  $`F_1 = \frac{2 \cdot \text{precision} \cdot \text{recall}}{\text{precision} + \text{recall}}`$.
  When the denominator is zero, scikit-learn's `zero_division=0` behavior
  reports zero.
- **Macro F1:** the arithmetic mean of the per-class F1 scores:
  $`F_{1}^{\text{macro}} = \frac{1}{C} \sum_{c=1}^{C} F_{1,c}`$.
  Every class contributes equally, regardless of its validation-set size
  ($`C`$ is the number of classes, $`F_{1,c}`$ the F1 of class $`c`$).
  This is the metric used for best-checkpoint selection.
- **Weighted F1:** the mean of per-class F1 scores weighted by each class's
  validation support. It is therefore more influenced by common classes.
- **Balanced accuracy:** the arithmetic mean of per-class recall:
  $`\text{balanced accuracy} = \frac{1}{C} \sum_{c=1}^{C} \text{recall}_c`$.
  It gives each class equal weight and is useful for imbalanced datasets.
- **Support:** the number of true validation examples for a class.

The confusion matrix uses rows for true classes and columns for predicted
classes. Diagonal entries are correct predictions. The compact confusion
summary reports each class's recall and its largest confusion destinations.


When TTA is enabled, validation loss is computed from the original image only,
while Top-1 accuracy, balanced accuracy, macro F1, weighted F1, and per-class
metrics use probabilities averaged over the original image and augmented
views. The log also reports original-view metrics and the change produced by
TTA.

Training metrics may be measured on an augmented or weighted-sampler stream.
They should not be compared directly with validation metrics unless they use
the same sampling and preprocessing scheme.

### Basic Fine-Tuning

```bash
genml-kit-train --model google/vit-base-patch16-224 \
                --dataset my-org/my-labeled-images \
                --label_column category \
                --epochs 5 \
                --batch_size 32 \
                --lr 3e-5 \
                --image_size 448
```

### With A Custom Classifier Head

Replace the default linear head with a custom MLP or attention-based
classifier:

```bash
# Freeze backbone, train only the custom head
genml-kit-train --model cls_model_wrapper:google/vit-base-patch16-224 \
                --dataset my-org/my-labeled-images \
                --label_column category \
                --model_arg classifier=mlp:hidden=512,dropout=0.3 \
                --freeze ".*\.(head|pool)"
```

### With LoRA (Parameter-Efficient Fine-Tuning)

Freeze the entire backbone and train only small low-rank adapter matrices.
Reduces trainable parameters by ~97% while often matching full fine-tuning:

```bash
genml-kit-train \
    --model cls_model_wrapper:facebook/dinov2-with-registers-large \
    --lora --lora_r 16 --lora_alpha 32 \
    --lora_target_modules "query,key,value" \
    --freeze "classifier\.(head|pool|encoder)" \
    --lr 3e-5 \
    --dataset my-org/my-labeled-images \
    --label_column category \
    --epochs 20
```

### From A Pre-Trained Checkpoint

Load encoder weights from a pre-training run (SimMIM, I-JEPA, or SupCon):

```bash
genml-kit-train --model convvit \
                --dataset my-org/my-labeled-images \
                --label_column category \
                --source_checkpoint ./checkpoints/convvit_simmim_latest.pt \
                --epochs 100
```

The backbone weights are loaded automatically; the classifier head is
reinitialised (different `num_classes`). Use `--state_load none` to avoid
carrying over old optimizer/scheduler states.

### Layer-Wise Learning Rate Decay (LLRD)

LLRD is a compromise between freezing the backbone and updating every layer at
the same speed. If layers are indexed from shallow 0 to deep $`L`$, a common
schedule is (introduced by Howard & Ruder, [Universal Language Model Fine-tuning](https://arxiv.org/abs/1801.06146), ACL 2018):

$$
\large
\text{lr}(\text{layer}) = \text{lr}_{\text{base}} \cdot d^{\,L - \text{layer}}
$$

where:

- $`\text{lr}_{\text{base}}`$ — the `--lr` you pass on the command line;
- $`d`$ — the decay factor (`--llrd_decay`), usually between 0.8 and 1.0;
- $`L`$ — the index of the deepest layer, $`\text{layer}`$ the current layer's
  index, both counted from 0 at the input side.

The deepest layer
receives the base rate while earlier layers receive smaller updates. Early
layers tend to represent general edges and textures; later layers are more
task-specific. This is a useful heuristic, not a law, so validate it on your
dataset. A very small decay factor can effectively freeze the shallow network.

### Cosine Annealing

`--scheduler CosineAnnealingLR` (the default in every example above) decays
the learning rate from its initial value $`\eta_{\max}`$ toward $`\eta_{\min}`$
(0 by default) following a half cosine wave over `--epochs` epochs $`T`$:

$$
\large
\eta(t) = \eta_{\min} + \frac{1}{2}\left(\eta_{\max} - \eta_{\min}\right)
\left(1 + \cos\left(\frac{t \pi}{T}\right)\right)
$$

where $`t`$ is the current epoch. The cosine shape spends more time at high
learning rates than a linear decay — useful early, when large steps explore
the loss surface — and slows down smoothly as it approaches zero, which lets
the model settle into a minimum instead of bouncing around it. Late-training
low learning rates are also what make the EMA teacher of DINO/BYOL
effective: the teacher tracks a student that changes gently.

### Mixup, Label Smoothing, And Imbalance

Mixup forms a virtual example from two training examples.
Reference: Zhang et al., [mixup: Beyond Empirical Risk Minimization](https://arxiv.org/abs/1710.09412), ICLR 2018.

$$
\large
\tilde{x} = \lambda x_i + (1 - \lambda) x_j
$$

$$
\large
\tilde{y} = \lambda y_i + (1 - \lambda) y_j
$$

where $`(x_i, y_i)`$ and $`(x_j, y_j)`$ are two randomly paired training
examples and the mixing coefficient $`\lambda`$ is drawn from the Beta
distribution, $`\lambda \sim \mathrm{Beta}(\alpha, \alpha)`$, with $`\alpha`$
set by `--mixup_alpha` (typical values 0.2–0.4). The parameter $`\alpha`$
shapes the distribution: small $`\alpha`$ puts most weight near 0 or 1,
large $`\alpha`$ near 1/2. The labels are probability vectors, not
class indices. Mixup smooths the decision boundary and can help on small
datasets, but strong Mixup can obscure fine-grained image details. Label
smoothing similarly replaces a one-hot label with a mostly-correct
distribution.

#### Focal Loss

Focal loss instead changes the emphasis of cross-entropy. Reference: Lin et
al., [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002), ICCV 2017.
With a hard label and $`p_t`$ the predicted probability of the correct class,
its basic form is

$$
\large
\mathcal{L}_{\text{focal}} = -(1 - p_t)^{\gamma} \log(p_t)
$$

where $`\gamma`$ (via `--loss_args '{"gamma": ...}'`) down-weights easy examples: when
$`p_t \to 1`$ the factor $`(1-p_t)^{\gamma} \to 0`$, so already-confident,
easy examples receive almost no weight and training focuses on hard ones.
With $`\gamma = 0`$ the factor is 1 for every example and the loss is plain
cross-entropy; typical values are 1–3, and 2.0 is the usual default.

genml_kit implements the general form (`CombinedFocalLoss` in
`genml_kit/losses/focal.py`), which composes cleanly with the other
tools in this section. Let $`q`$ be the target distribution from the previous
step (one-hot, Mixup-blended, or label-smoothed), $`w_c`$ the per-class
multipliers (`--class_multipliers`), and $`p`$ the model's softmax output.
The loss for one example is

$$
\large
\mathcal{L} = \left(1 - p_t\right)^{\gamma} \sum_{c=1}^{C} w_c\, q_c
\left(-\log p_c\right), \qquad p_t = \sum_{c=1}^{C} q_c\, p_c
$$

Two details matter for the reader combining these flags:

- **Unified $`p_t`.** The modulating factor uses the *expected* probability
  under the target distribution, $`p_t = \sum_c q_c p_c`$, rather than the
  probability of a single class. For a one-hot target $`q`$ this reduces
  exactly to the basic form above ($`p_t`$ becomes the correct-class
  probability); with Mixup soft targets it blends the two classes'
  probabilities, giving one smooth modulation per example instead of an
  arbitrary single-class choice.
- **One modulation per example.** The factor $`(1-p_t)^{\gamma}`$ scales the
  whole weighted cross-entropy sum — weights and soft labels act *inside* the
  sum, focal acts *outside* it. Each term is therefore attenuated
  proportionally to how confident the model already is about that example, no
  matter which class the probability mass sits on.

Use these tools deliberately; combining every regularizer is not
automatically better.

### Hyperparameter Guidance

| Scenario | `--lr` | `--epochs` | `--batch_size` | Notes |
|---|---|---|---|---|
| Large dataset (>10k images) | 3e-5 | 20–50 | 32–64 | Standard fine-tuning |
| Small dataset (<1k images) | 1e-5 | 50–100 | 16–32 | Consider LoRA, stronger augmentation |
| From pre-trained checkpoint | 3e-5 | 50–100 | 32 | Lower LR than from scratch |
| Custom classifier head only | 1e-3 | 50–200 | 32 | Higher LR since only head trains |

**Tips:**
- Start with `--lr 3e-5` for full fine-tuning, `1e-3` for head-only training.
- Use `--mixup_alpha 0.2` for small datasets — it helps prevent overfitting.
- `--loss_args '{"gamma": 2.0}'` down-weights easy examples, useful when
  classes are imbalanced.
- `--class_multipliers "rare_class=3.0"` increases the loss weight for
  safety-critical or otherwise priority classes.

### Reproducibility

RNG streams are always seeded.  By default the seed comes from the
`GENML_KIT_SEED` environment variable (default `42`), and it drives
DataLoader shuffling, per-worker augmentation randomness, mixup, and
dropout — the user not caring about determinism does not mean the
toolkit runs on entropy.  To compare hyperparameters under a different
random draw, change `GENML_KIT_SEED` or pass an explicit seed.

Two levels of determinism:

| Invocation | Seeding | CUDA kernels | Reproducible? |
|---|---|---|---|
| default (no `--seed`) | `GENML_KIT_SEED` or 42 | `cudnn.benchmark=True` (fastest algorithms) | RNG streams identical; arithmetic not bit-exact |
| `--seed 42` | 42 everywhere | cuDNN deterministic mode, benchmark off, deterministic algorithms | bit-exact |

Passing an explicit `--seed` therefore buys full determinism at a
throughput cost (typically 10–20%); omit it for maximum speed:

```bash
genml-kit-train --model convvit --seed 42 ...
```

The train/val split and the XGBoost stage follow the user's seed: with
`--seed` given they are pinned; without it they draw fresh entropy each
run (HF `train_test_split(seed=None)`, XGBoost `random_state=None`).

Two caveats:

- Some CUDA ops have no deterministic kernel.  Instead of aborting a
  long-running job, genml_kit logs a warning and proceeds (the op falls
  back to a non-deterministic kernel).
- `float16` AMP with `GradScaler` involves non-associative reductions
  that can still differ run-to-run; use `bfloat16` (default on
  Ampere+) or disable AMP for strictly repeatable arithmetic.

Checkpointing is atomic: `_latest.pt` is written to a temporary file
and renamed into place, so an interrupted run never leaves a truncated
resume point.

### Fine-Tuning CLI Reference

| Argument | Default | Description |
|---|---|---|
| `--model` | `google/vit-base-patch16-224` | HuggingFace model name, local path, or custom model (e.g. `convvit`, `timm:<name>`). |
| `--dataset` | required | HuggingFace dataset name or `imagefolder/PATH` for local data. |
| `--image_column` | auto-detected | Explicit HF image column name. |
| `--label_column` | auto-detected | Explicit HF label column name. |
| `--image_size` | `448` | Augmentation crop size (processor handles final resize). |
| `--epochs` | `5` | Number of training epochs. |
| `--batch_size` | `32` | Batch size. |
| `--lr` | `3e-5` | Peak learning rate. |
| `--weight_decay` | `0.01` | Weight decay. |
| `--loss` | `None` | Loss spec: registered name (`focal`, `supcon`, or the RL bundles `ppo`/`sac`/`dqn`) or a path/URL to a `.py` script defining `build_loss(**kwargs)`. `None` (default) selects the method's built-in loss. |
| `--loss_args` | `{}` | JSON kwargs forwarded to the loss builder. Example: `'{"gamma": 1.0, "label_smoothing": 0.1}'`. |
| `--class_multipliers` | `""` | Per-class priority multipliers. Example: `"cat=3.0,dog=1.0"`. |
| `--sampler` | `none` | Training sampler: `none` (shuffle), `weighted` (WeightedRandomSampler for class imbalance), or `balanced` (equal samples per class per batch; batch_size need not divide evenly). |
| `--sampler_weights` | `frequency` | Weight mode for `--sampler weighted`: `frequency` (inverse-freq), `multipliers` (--class_multipliers), or `combined` (freq × multipliers). |
| `--mixup_alpha` | `0.0` | Mixup alpha (`0` = disabled; recommended: `0.2`). |
| `--seed` | `None` | Explicit RNG seed: pins the train/val split and the XGBoost stage and enables full determinism (cuDNN deterministic kernels, benchmark off). Omit (default) to keep runs fast — shuffling/mixup/dropout are still seeded from `$GENML_KIT_SEED` (default 42). See [Reproducibility](#reproducibility). |
| `--grad_accum_steps` | `1` | Gradient accumulation steps (effective batch = batch_size × steps). |
| `--amp_dtype` | `None` | Mixed precision: `float16` or `bfloat16`. |
| `--device` | auto-detect | Device: `cpu`, `cuda`, or `cuda:INDEX`. |
| `--lr_group` | `None` | Per-parameter-group learning rates (repeatable). Format: `"REGEX=LR"`. |
| `--llrd_decay` | `None` | Layer-wise LR decay factor per depth level. Example: `--llrd_decay 0.85`. |
| `--checkpoint` | `genml_kit` | Checkpoint base path (`_latest.pt` / `_best.pt` appended). |
| `--log_every` | `20` | Log every N steps. |
| `--grad_monitor` | `-1` | Log gradient statistics every N steps. See [Gradient Monitor](#gradient-monitor). |
| `--norm_history` | `0` | Keep last N norm snapshots for trend analysis. |
| `--trend_top_n` | `10` | Show top N params in trend table. `0` = show all. |
| `--grad_clip` | `1.0` | Max gradient norm for clipping. `0` disables. |
| `--save_every` | `500` | Save checkpoint every N optimizer steps. 0 disables. |
| `--num_workers` | `2` | DataLoader worker processes. |
| `--log_level` | `INFO` | Logging level. |
| `--log_targets` | `STDERR` | Comma-separated log destinations. `STDERR` logs to standard error; any other entry is a log file path (appended). Example: `STDERR,/tmp/train.log` logs to both. |
| `--log_dir` | `None` | TensorBoard log directory (default: `<checkpoint_dir>/logs`). Requires the `tensorboard` package (installed by the `dev` extra). |
| `--cache_dir` | `None` | HuggingFace cache directory. |
| `--remote_checkpoint` | `None` | Remote URI for checkpoint sync (`gs://BUCKET/PREFIX`, `r2://BUCKET/PREFIX`, or `s3://BUCKET/PREFIX`). |
| `--source_checkpoint` | `None` | Path to source checkpoint to absorb parameters from. |
| `--param_rename` | `None` | Regex-based key rename patterns (`SEARCH;REPLACE`). |
| `--freeze` | `None` | Regex patterns for parameters to keep trainable. All others frozen. |
| `--lora` | `False` | Enable LoRA via PEFT. Requires `pip install "genml_kit[lora]"`. |
| `--lora_r` | `8` | LoRA rank. |
| `--lora_alpha` | `16` | LoRA alpha (scaling = `alpha / r`). |
| `--lora_dropout` | `0.0` | Dropout on LoRA layers. |
| `--lora_target_modules` | `None` | Comma-separated module names for LoRA (e.g. `"query,key,value"`). |
| `--optimizer` | `AdamW` | `torch.optim` optimizer class name or `.py` script path. |
| `--opt_arg` | `{}` | Extra optimizer kwargs (repeatable). |
| `--scheduler` | `None` | `torch.optim.lr_scheduler` class name or `.py` script path. |
| `--sched_arg` | `{}` | Extra scheduler kwargs (repeatable). |
| `--state_save` | `opt,sched,amp` | States to save: `opt`, `sched`, `amp`, `none`. |
| `--state_load` | `opt,sched,amp` | States to restore on resume. |
| `--xgboost_model` | `None` | Output path for XGBoost model (trains after PyTorch). |
| `--xgb_*` | various | XGBoost hyperparameters (see `--help` for full list). |
| `--model_arg` | `{}` | Override model configuration (repeatable). |
| `--proc_arg` | `{}` | Override processor configuration (repeatable). |
| `--train_augmentation_script` | `None` | Custom augmentation script. Must define `create_train_transform()`. |
| `--grad_checkpoint` | `False` | Enable gradient checkpointing. Reduces activation memory by ~40-50% at the cost of ~25-35% more compute per step. Enables larger batch sizes. See [Gradient Checkpointing](#gradient-checkpointing). |
| `--tta` | `None` | Test-Time Augmentation. `default` uses built-in 4-view transform (identity + flips). A path/URL loads an external script defining `create_tta_transform()`. Omit to disable. |

Training automatically resumes from an existing `_latest.pt` or `_best.pt`
checkpoint if one exists at the `--checkpoint` path.

### Remote Checkpoint Sync (GCS / R2 / S3)

`--remote_checkpoint` uploads each saved checkpoint to cloud storage.
Requires `pip install "genml_kit[s3]"` for `s3://` and `r2://` URIs (both
use boto3), or `genml_kit[gcs]` for `gs://`. The unified `genml-kit-train`
CLI accepts the flag for every pipeline/method configuration.

The sync also works in reverse at startup: before auto-resume, any missing
`_latest.pt` / `_best.pt` is downloaded from the remote prefix (latest is
tried first, matching resume precedence). A locally present checkpoint is
never overwritten — the local copy always wins — and connection or
credential failures degrade to a warning instead of aborting startup.

**AWS S3** — credentials come from the standard environment variables:

```bash
%env AWS_ACCESS_KEY_ID=AKIA...
%env AWS_SECRET_ACCESS_KEY=...
%env AWS_SESSION_TOKEN=...        # only for temporary (STS/SSO) credentials
%env AWS_DEFAULT_REGION=us-east-1

--remote_checkpoint s3://my-bucket/genml_kit/convvit_ijepa
```

Permanent IAM-user keys need only the access/secret pair; the session
token is picked up automatically when present. When the key variables
are unset, boto3's default credential chain applies (IAM instance role,
`~/.aws/credentials`, SSO cache).

**Cloudflare R2** — the S3-compatible API with R2 credentials:

```bash
%env R2_ENDPOINT_URL=https://<account_id>.r2.cloudflarestorage.com
%env R2_ACCESS_KEY_ID=...
%env R2_SECRET_ACCESS_KEY=...

--remote_checkpoint r2://my-bucket/genml_kit/convvit_ijepa
```

### LoRA Details

Low-Rank Adaptation freezes the pre-trained backbone and injects small
trainable low-rank matrices into attention layers.
Reference: Hu et al., [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685), ICLR 2022.

Instead of updating a weight matrix $`W`$ directly, LoRA keeps it frozen and
adds a low-rank product to its output:

$$
\large
W' = W + \Delta W, \qquad \Delta W = \frac{\alpha}{r} \cdot B A
$$

where:

- $`W \in \mathbb{R}^{d \times k}`$ — the frozen pre-trained weight matrix;
- $`A \in \mathbb{R}^{r \times k}`$, $`B \in \mathbb{R}^{d \times r}`$ — the two
  small trainable matrices ($`BA`$ is a matrix product);
- $`r`$ — the rank (`--lora_r`), much smaller than $`d`$ or $`k`$, which is
  what makes $`\Delta W`$ cheap: instead of $`d \times k`$ weights it trains
  only $`r(d + k)`$;
- $`\alpha`$ — the scale factor (`--lora_alpha`); the effective step size of
  the update is $`\alpha / r`$.

Only $`A`$, $`B`$ (and the classifier head) receive gradients.

| r | alpha | alpha/r | Use case |
|---|---|---|---|
| 8 | 16 | 2.0 | Conservative, very few params |
| 16 | 32 | 2.0 | Good default for medium datasets |
| 16 | 64 | 4.0 | Larger updates (helpful under bfloat16) |
| 32 | 64 | 2.0 | More capacity, same scaling |

LoRA can be combined with a custom classifier:

```bash
genml-kit-train \
    --model cls_model_wrapper:facebook/dinov2-with-registers-large \
    --model_arg classifier=cls_attention:num_encoder_layers=2 \
    --lora --lora_r 16 --lora_alpha 32 \
    --freeze 'classifier\.(head|pool|encoder)' \
    --lr_group 'backbone.*=1e-5' 'classifier.*=3e-4' \
    ...
```

### Cross-Dataset Resume

Switch from one dataset to another while keeping backbone weights:

```bash
genml-kit-train --model facebook/convnextv2-base-22k-224 \
                --dataset my-org/other-labeled-images \
                --label_column category \
                --checkpoint genml_kit \
                --state_load none \
                --epochs 10 \
                --batch_size 16 \
                --lr 3e-5 \
                --mixup_alpha 0.2 \
                --amp_dtype bfloat16
```

Backbone weights load via `strict=False`; the classifier head (different
`num_classes`) is reinitialised. `--state_load none` prevents carrying over
old optimizer/scheduler states.

### Gradient Checkpointing

Training deep vision transformers is memory-intensive: the attention maps for
EVA02-base at 448×448 resolution consume ~4.6 GB across 12 layers, which
limits the maximum batch size even on a 24 GB GPU. Gradient checkpointing
trades compute for memory by discarding intermediate activations during the
forward pass and recomputing them during the backward pass.

```bash
genml-kit-train \
    --model timm:eva02_base_patch14_448.mim_in22k_ft_in22k_in1k \
    --batch_size 32 \
    --grad_accum_steps 2 \
    --grad_checkpoint \
    --amp_dtype bfloat16 \
    ...
```

**What it does:** Each transformer block stores only its input and output.
The intra-block activations (attention maps, FFN intermediates) are
recomputed on-demand during the backward pass. This reduces peak activation
memory by ~40-50%.

**What it costs:** The recomputation adds ~25-35% more compute per training
step. Whether this translates to a net throughput gain or loss depends on
your GPU:
- **Memory-bound GPU** (VRAM-limited, compute underutilized): larger batches
  fill the GPU's compute capacity → net throughput **increase** of ~50-60%.
- **Compute-saturated GPU** (TFLOPS at 100%): every extra FLOP adds to step
  time → net throughput **decrease** of ~20-25%, though the larger batch can
  still improve training stability and final accuracy.

**Backend support:** Enabled automatically for all model backends — timm,
HuggingFace, ConvViT, and UVito — via native APIs or per-block
`torch.utils.checkpoint` in the transformer loops.

---

## Inference Guide

Run inference on individual images:

```bash
genml-kit-infer --model facebook/convnextv2-base-22k-224 \
                --checkpoint genml_kit_best.pt \
                path/to/image.jpg path/to/other_image.png
```

Output is JSON with per-class probabilities:

```json
{
  "source": "image.jpg",
  "predictions": [
    {"label": "golden_retriever", "probability": 0.435},
    {"label": "labrador", "probability": 0.281}
  ]
}
```

### XGBoost And Test-Time Augmentation

The neural classifier makes decisions through its head. The XGBoost option
takes the encoder representation $`h = f_\theta(x)`$ instead and fits an
ensemble of decision trees to those vectors.
Reference: Chen & Guestrin, [XGBoost: A Scalable Tree Boosting System](https://arxiv.org/abs/1603.02754), KDD 2016.
A tree partitions feature space with rules such as $`h_{37} < t`$ (comparing
the 37th feature of $`h`$ against a learned threshold $`t`$); boosting adds
trees sequentially so each new tree focuses on errors left by previous trees.
This can work well when the dataset is small
and the representation is already useful, but it is not guaranteed to beat the
neural head. Use validation data to choose tree depth, number of rounds, and
ensemble weight.

Test-time augmentation (TTA) runs the same image through several plausible
views and averages the probability vectors:

$$
\large
\bar{p}(y \mid x) = \frac{1}{K} \sum_{k=1}^{K} p(y \mid T_k(x))
$$

where $`T_1, \dots, T_K`$ are the $`K`$ test-time transformations (e.g. the
original image plus flipped copies) and $`\bar{p}`$ is their averaged
probability vector: each transformed copy is classified independently, and the
final prediction averages the per-view class probabilities. The
transformations should preserve the label semantics. Horizontal flips
are usually safer than orientation-specific crops; verify that an augmentation
does not erase the cue that distinguishes the classes.

### Inference CLI Reference

| Flag | Default | Description |
|---|---|---|
| `--model` | (required) | HuggingFace model name or custom model. |
| `--checkpoint` | (required) | Path to state dict or wrapped checkpoint. |
| `--top_k` | `None` | Show top-K predictions; omit for all classes. |
| `--output` | `None` | Write JSON results to file. |
| `--device` | `None` | Force PyTorch device (`cuda`, `cpu`). Auto-detected if omitted. |
| `--cache_dir` | `None` | HuggingFace cache directory. |
| `--xgboost_model` | `None` | XGBoost model path. If provided, runs XGBoost alongside PyTorch. |

Wrapped checkpoints (containing `model_state_dict` and metadata) are preferred
over raw state dictionaries, which produce a metadata warning.

### XGBoost Inference

When `--xgboost_model` is provided, the output includes both predictions:

```json
{
  "source": "image.jpg",
  "predictions": [
    {"label": "golden_retriever", "probability": 0.435}
  ],
  "xgboost_predictions": [
    {"label": "golden_retriever", "probability": 0.612}
  ]
}
```

---

## Tips & Pitfalls

### `--checkpoint` And `--source_checkpoint` Are Different

`--checkpoint` names the output prefix used for saving and resuming the current
training run. `--source_checkpoint` imports weights from another run before
the new optimizer is created. When moving a SupCon encoder into classification,
use:

```bash
genml-kit-train \
    --checkpoint /content/eva02_finetune \
    --source_checkpoint /content/eva02_supcon_latest.pt \
    --param_rename 'encoder\\.model\\.(.*);model.$1' \
    ...
```

The rename is needed because the pre-training wrapper stores backbone keys
under `encoder.model.*`, whereas the fine-tuning model stores them under
`model.*`. The SupCon `projection.*` keys are expected to remain unused, and
new classifier-head keys are expected to be missing from the source checkpoint.
Those messages indicate a successful partial transfer if the backbone keys
are matched.

### Interpreting A SupCon Plateau

SupCon loss is not expected to approach zero. If an anchor has k-1 positives and
positives become much more similar than negatives, a useful reference value is
approximately `-log(k-1)`. This is only a diagnostic approximation: class
counts may vary, the sampler may repeat examples, and temperature affects
optimization. Judge the loss together with embedding quality and downstream
validation metrics. A plateau near this reference can mean convergence; a
plateau well above it can mean too few positives, a learning-rate problem, or
labels that are not being passed correctly.

### A Practical Debugging Order

1. Confirm that images and labels are valid and that each SupCon batch has
   repeated classes.
2. Confirm that the loss changes when the learning rate changes and inspect
   `--grad_monitor` for zero or exploding gradients.
3. Check whether pre-trained weights loaded by reading the alignment report,
   not just the final epoch number.
4. Compare against a simple head-only or full-fine-tuning baseline before
   adding LLRD, LoRA, Mixup, focal loss, and class multipliers together.
5. Select the checkpoint using a validation metric appropriate to the
   objective, rather than training loss alone.

### Low GPU Utilization (`util=0` In The Training Log)

The periodic `GPU:` line reports memory in MiB and SM utilization as a
percentage sampled at log time (`mem=allocated/totalMiB (pct) res=…MiB
util=…%`). `mem` counts live tensors, `res` is the caching allocator's
pool, and `util` is one instant — not an average — so an isolated 0 is
normal. A *persistent* pattern of `util=0` at every log line, with
`mem`/`res` byte-identical between samples, means the GPU finished the
last batch and is idle waiting for the next one: the run is **input
bound**, not compute bound. The throughput line (`img/s=`) is the ground
truth for step time.

Common causes and fixes, in order of likelihood:

- **Too few DataLoader workers.** JPEG decode + augmentation is CPU work;
  with `--num_workers 4` a 4090 can easily outrun the pipeline at
  448². Check `nproc` on the machine and raise `--num_workers` toward
  the core count — `img/s` should scale until the GPU saturates
  (a `util` that now sits consistently above ~90% means you are done).
- **Expensive decode.** Full-resolution JPEGs dominate decode cost.
  Pre-resizing the dataset offline (e.g. ~512px long side) cuts decode
  work several-fold and often doubles throughput on small VRAM budgets.
- **Slow storage.** A network volume (S3/R2 mount, NFS) adds latency per
  image; prefer local NVMe for the dataset.
- **`--seed`.** It enables cuDNN deterministic mode and disables
  benchmarking; expect slower steps by design.  Omit it when you do not
  need bit-exact runs.

Fine-tuning runs on frozen-backbone models need very little GPU memory
(`mem` is typically a few hundred MiB), so a low `mem` number is normal
there and not a sign of a problem. Loader prefetch (`prefetch_factor=4`
per worker) and `cudnn.benchmark=True` (default; disabled by an explicit
`--seed`) already help absorb jitter and speed up the GPU side.

---

## Gradient Monitor

`--grad_monitor N` logs a per-parameter gradient report every *N* training
steps. This helps diagnose training instability (exploding / vanishing
gradients, imbalanced parameter updates) before it shows up in the loss.

### Summary Line

```
[Step 29400] Gradient Report: 202 params | grad_rms: mean=5.68e-01 max=3.41e+00 min=1.66e-05 | grad/param: mean=2.94e-01
```

All norms are **RMS** (root mean square): the $`L_2`$ norm of the tensor
divided by $`\sqrt{N}`$, where $`N`$ is the number of elements (`numel`):

$$
\large
\mathrm{RMS}(v) = \frac{\lVert v \rVert_2}{\sqrt{N}} = \sqrt{\frac{1}{N} \sum_{n=1}^{N} v_n^2}
$$

Here $`v`$ is the tensor (a parameter or its gradient) flattened to $`N`$
values $`v_1, \dots, v_N`$, and $`\lVert v \rVert_2 = \sqrt{\sum_n v_n^2}`$ is
its Euclidean norm. Dividing by $`\sqrt{N}`$ makes the measure independent of
tensor shape and directly comparable across parameters of different sizes.

| Field | Meaning |
|---|---|
| `params` | Total number of trainable parameters. |
| `grad_rms: mean/max/min` | RMS of the gradient tensor for each parameter, then aggregated. `max` is the single most aggressive gradient — the one most likely to cause instability. |
| `grad/param: mean` | Average gradient-to-parameter ratio. Healthy: < 0.1. Concerning: > 1.0. |

### Per-Parameter Columns

| Column | Symbol | What to look for |
|---|---|---|
| **g_rms** | $`\mathrm{RMS}(\nabla_W \mathcal{L})`$ | Compare across params. One param with g_rms 100× higher is a problem. |
| **p_rms** | $`\mathrm{RMS}(W)`$ | Per-element scale context. With `std=0.02` init, expect ~0.02. |
| **g/p** | $`\lVert \nabla_W \mathcal{L} \rVert_2 / (\lVert W \rVert_2 + \varepsilon)`$ | **Most useful column.** Healthy: < 0.1. Concerning: > 1.0 (update overshoots). Dangerous: > 5.0. |
| **g_max** | $`\max_n \lvert \partial \mathcal{L} / \partial W_n \rvert`$ | Highlights individual neurons with extreme gradients. |
| **sparse** | `% zero` | High sparsity (> 50%) = most neurons not receiving signal. |
| **status** | | `OK` = healthy. `STL` = stalled. `OVF` = exploding. `IMB` = imbalanced. `GPR` = high g/p ratio. |

In this table $`W`$ denotes one parameter tensor, $`\nabla_W \mathcal{L}`$ its
gradient (the vector of partial derivatives of the loss with respect to each
element of $`W`$), and $`\varepsilon`$ a tiny constant that only prevents
division by zero for parameters that are exactly zero.

### Reading The Report

**Healthy training:** g/p < 0.1, `grad/param: mean` in 0.01–0.1 range, g_rms within ~10× across layers.

**Exploding gradients:** One or more params with `OVF`, g/p > 5.0, `grad_rms: max` >> `mean`. Fix: lower LR, add `--grad_clip 1.0`, or use warmup.

**Vanishing gradients:** Many params with `STL`, g_rms near 1e-7, high sparsity. Fix: increase LR, check for dead neurons.

**Imbalanced updates:** Some params `IMB`, large g/p disparity between layers. Fix: use `--lr_group` for different rates, or freeze the dominant component.

### Norm Trend History

`--norm_history N` (requires `--grad_monitor`) keeps the last N snapshots
per parameter. A trend summary table is appended to each report showing
direction (`UP`/`DOWN`/`---`), percentage change, and min/max values.

---

## Custom Losses

Losses are resolved through the `LOSSES` registry
(`genml_kit.losses.registry.load_loss`).  A loss spec is either a
registered name (`focal`, `supcon`, or the RL bundles `ppo`, `sac`,
`dqn`) or a path/URL to a `.py` script defining `build_loss(**kwargs)`:

```python
# my_loss.py
import torch.nn as nn

class MyLoss(nn.Module):
  def __init__(self, alpha=1.0):
    super().__init__()
    self.alpha = alpha

  def forward(self, logits, targets):
    return self.alpha * nn.functional.cross_entropy(logits, targets)

def build_loss(alpha=1.0, **kwargs):
  """Protocol: return a callable (logits, targets) -> tensor."""
  return MyLoss(alpha=alpha)
```

```bash
genml-kit-train --model google/vit-base-patch16-224 \
  --loss ./my_loss.py --loss_args '{"alpha": 2.0}' ...
```

**RL losses** work the same way but return a *bundle* (dict of
callables); the script's entries are merged over the method's built-in
bundle (the method passes its own name as the merge base: `ppo`, `sac`,
or `dqn`), so a script may override just one loss and inherit the rest:

```python
# my_ppo_entropy.py
def build_loss(**kwargs):
  return {"entropy": lambda ent: ent.mean() * 2.0}  # override only entropy
```

```bash
genml-kit-train --method ppo --loss ./my_ppo_entropy.py ...
```

The built-in bundles are merged over (`ppo`: `policy`/`value`/`entropy`;
`sac`: `q`/`policy`/`alpha`; `dqn`: `td`).  The `--loss_args` JSON
kwargs are forwarded to `build_loss(**kwargs)`.

## Custom Models

genml_kit supports any HuggingFace `AutoModelForImageClassification` model,
any timm model via `timm:<name>`, and custom architectures registered in
`genml_kit.models`.

### Built-In Custom Models

| Name | Description | `--model` value |
|---|---|---|
| timm | Any model from [timm](https://github.com/huggingface/pytorch-image-models) | `timm:<model_name>` |
| ConvViT | Multi-block conv stem + ViT encoder with CLS-guided attention pooling | `convvit` |
| UVito | Frozen SMP encoder + learnable patch projection + Transformer encoder | `uvito` |
| ClsModelWrapper | HuggingFace backbone + custom classifier head | `cls_model_wrapper:<hf_name>` |
| ContrastiveEncoder | Backbone + projection head for contrastive pre-training | `contrastive_encoder:<hf_name>` |

### Adding A Custom Model

1. Create `genml_kit/models/{name}/` with `model.py`, `processor.py`,
   `loader.py`, and `__init__.py`.
2. Add the import to `genml_kit/models/__init__.py`.
3. The model must expose `.forward(pixel_values=images)` → object with `.logits`,
   and `config.id2label` / `config.label2id`.
4. CLI overrides via `--model_arg KEY=VALUE` are forwarded to the loader.

### External Model Scripts

Instead of registering a model inside the source tree, you can point
`--model` at a `.py` file (or HTTP/HTTPS URL).  The script must define
`create_model`:

```python
# my_model.py
import torch.nn as nn

class Tiny(nn.Module):
  def __init__(self, num_labels, hidden=128):
    super().__init__()
    self.net = nn.Sequential(nn.Linear(784, hidden), nn.ReLU(),
                             nn.Linear(hidden, num_labels))

  def forward(self, pixel_values):
    return pixel_values

def create_model(num_labels, image_size=224, id2label=None, label2id=None,
                 checkpoint_path=None, device="cpu", **kwargs):
  """Protocol: return a torch.nn.Module."""
  return Tiny(num_labels=num_labels, hidden=kwargs.get("hidden", 128))
```

```bash
genml-kit-train --model ./my_model.py --loss focal ...
```

The script is **fully responsible for weight loading**: when
`checkpoint_path` is provided (inference), call
`genml_kit.io.checkpointing.load_checkpoint_weights(checkpoint_path, model, ...)`
inside `create_model`, mirroring the ConvViT/UViTO loaders.  In training
`checkpoint_path` is `None` (random init) and the global `--source_checkpoint`
/ `--resume` paths apply uniformly after `create_model` returns.

An optional `create_processor(image_size=224, **kwargs)` may be defined
for the processor; it must return an object exposing `image_mean` /
`image_std`.

### Custom Classifiers

`ClsModelWrapper` lets you replace the default HF classification head:

```python
class Classifier(nn.Module):
    def __init__(self, num_labels, hidden_size, **kwargs):
        super().__init__()
        self.head = nn.Linear(hidden_size, num_labels)

    def forward(self, hidden_states):          # (B, N, D) tensor
        features = hidden_states[:, 0]         # CLS token
        return self.head(features)

    def extract_features(self, hidden_states): # (B, N, D) → (B, D)
        return hidden_states[:, 0]
```

The `extract_features` method is used by `--xgboost_model` for XGBoost
training on backbone features.

The head is selected through `--model_arg`, not a dedicated flag:

```bash
--model cls_model_wrapper:<hf_name> --model_arg classifier=mlp:hidden=512,dropout=0.3
```

The `classifier` value is an inline spec: the head name (a registered
built-in — `mlp`, `cls_attention` — or a path to a `.py` file defining a
`Classifier` class), optionally followed by `:key=value,...` kwargs. Values
are parsed as int/float/bool, and bracketed lists such as
`cls_slice=(0, 1)` work. A head is required: `cls_model_wrapper` fails fast
with a clear message when `classifier` is missing. Models that embed their
own head (`convvit`, `uvito`, `timm:<name>`) do not accept this kwarg — for
a custom head on a HuggingFace backbone use `cls_model_wrapper` instead.
`genml-kit-infer` needs the same `--model_arg classifier=...` to rebuild the
head before loading the checkpoint weights.

---

## Reinforcement Learning

genml_kit ships a self-contained reinforcement-learning (RL) toolkit built on
Gymnasium.  The companion document
[`rl/README.md`](rl/README.md) is a from-the-ground-up *mathematics* tutorial
(Bellman equations, TD, Q-learning, GAE, PPO, SAC); this section is the
practical *usage* guide.

### Installation

RL support is an optional extra:

```
pip install 'genml_kit[rl]'
```

### Supported algorithms

| `--method` | Algorithm | Action space | Key files |
|---|---|---|---|
| `dqn` | Double-DQN with epsilon-greedy, optional dueling + n-step + PER | Discrete | `genml_kit/methods/rl_dqn.py` |
| `ppo` | On-policy PPO with GAE, clipped surrogate, entropy bonus | Discrete or continuous | `genml_kit/methods/rl_ppo.py` |
| `sac` | Soft Actor-Critic, twin critics, reparameterized policy, auto-alpha | Continuous only | `genml_kit/methods/rl_sac.py` |

### Quick start

```
# DQN on CartPole (discrete)
genml-kit-train --pipeline rl --method dqn --env_id CartPole-v1 \
    --epochs 20 --batch_size 64

# PPO on CartPole (discrete)
genml-kit-train --pipeline rl --method ppo --env_id CartPole-v1 \
    --epochs 20 --ppo_rollout_len 2048

# SAC on a continuous env (e.g. LunarLanderContinuous-v3)
genml-kit-train --pipeline rl --method sac --env_id LunarLanderContinuous-v3 \
    --epochs 20
```

All RL methods use `eval_return` (mean undiscounted return over
`--eval_episodes` episodes, default 5) as the checkpoint-selection metric.
The trainer bypasses the DataLoader: it steps the environment directly and
samples from a replay buffer (DQN/SAC) or re-uses a rollout buffer (PPO).

### Custom RL models

Every RL method resolves its network through the same model registry as the
rest of the framework: `--model` selects a registered factory and
`--model_arg key=value` forwards constructor kwargs.

```
# DQN: dueling Q-head (also the default for --dqn_dueling)
genml-kit-train --pipeline rl --method dqn --model rl/qnet_dueling

# DQN: configure the MLP width via --model_arg
genml-kit-train --pipeline rl --method dqn \
    --model rl/qnet --model_arg hidden_dims=[512,256]

# PPO: wider actor-critic trunk
genml-kit-train --pipeline rl --method ppo \
    --model rl/actor_critic --model_arg hidden_dims=[64,64]

# SAC: critic head width (the actor defaults to rl/actor_critic)
genml-kit-train --pipeline rl --method sac --env_id Pendulum-v1 \
    --model rl/sac_critic --model_arg hidden_dims=[256,256]
```

Registered RL factories (auto-discovered; see `MODELS.list_names()`):

| Name | Returns | Notes |
|---|---|---|
| `rl/qnet` | `QNetwork` (flat Q-head) | DQN default |
| `rl/qnet_dueling` | `QNetwork` (dueling Q-head) | `--dqn_dueling` shortcut |
| `rl/actor_critic` | `ActorCritic` (shared backbone) | PPO default |
| `rl/sac_critic` | `SACCritic` (single critic head) | SAC default |

The methods pass a `SpaceSpec` (obs/action sizes derived from the env) to
the factory, so a custom registered factory can accept the standard scalar
kwargs (`obs_dim`, `n_actions`, `action_dim`, `hidden_dims`, ...) or the
`space=` spec directly.  SAC additionally supports a *factory-return
contract*: a factory may return a plain dict `{"actor": ..., "critic": ...}`
to supply both networks (the method composes the twin critic via a deep
copy).  As elsewhere in the framework there is **no** model-validation
layer: an incompatible network fails loudly at first use (`act()` /
`forward()`), and an unknown `--model` name falls through `load_model` to
Hugging Face and fails at load time.

### Evaluation videos

Pass `--record_eval_video` to record **one video per evaluation episode**
during validation (all RL methods: DQN / PPO / SAC).
`<checkpoint_dir>/videos/` as `eval_episode_000.mp4` etc., or as
`eval_episode_000.gif` when the MP4 writer is unavailable (e.g. no
`imageio[ffmpeg]`).  Recording is **only** attempted when the environment can
produce `rgb_array` frames: scripted/custom environments without a renderer
are detected at validation time and skipped without crashing.

### How a PPO training epoch works

PPO is *on-policy*: every update must be computed from transitions
collected by the **current** policy.  To keep that guarantee, the rollout
buffer is **refilled from scratch each epoch** — nothing survives from the
previous epoch.

**One training epoch** (the `=== Epoch N/500 ===` line in the logs) is
exactly:

```
1.  rollout.reset()                        # clear the on-policy buffer
2.  collect rollout_len env steps          # --ppo_rollout_len (default 2048)
3.  GAE (gamma, lam) over the rollout      # advantages + returns
4.  normalize advantages once              # over the whole rollout
5.  ppo_epochs SGD passes over it          # --ppo_epochs (default 4)
    each pass: rollout_len / mini_batch_size minibatches
6.  next epoch -> back to 1
```

So the arithmetic of a run is:

| Quantity | Formula | Default (Pendulum) |
|---|---|---|
| env steps per epoch | `rollout_len` | 2048 |
| gradient steps per epoch | `ppo_epochs * (rollout_len / mini_batch_size)` | `4 * 2048/64 = 128` |
| episodes per epoch | `rollout_len / episode_len` | `2048/200 ≈ 10` |
| total env steps | `epochs * rollout_len` | 500 × 2048 = 1.02M |

The **two "epochs"** that can be confused:

- **Training epoch** (`--epochs`) — one rollout + its updates.  This is
  what the log lines and checkpoints count.
- **PPO SGD epoch** (`--ppo_epochs`) — the number of *passes over the
  same* rollout during the update phase.  Raising it increases data
  reuse of a single rollout; lowering it (or raising `rollout_len`)
  increases data freshness.

The buffer itself is wiped cheaply: `rollout.reset()` just rewrites the
pointer, so the tensor storage is reused and the per-epoch cost is one
`rollout_len` collection, not an allocation.  On resume, the rollout
buffer state is saved/restored with the checkpoint and the previous
epoch's last observation carries over, so an interrupted run continues
mid-environment rather than restarting the episode.

Changing `--ppo_rollout_len` therefore changes the *epoch length in
steps*, not the memory footprint (memory is `O(rollout_len)` but tiny —
~115 KB for Pendulum at 2048) and not the total step budget
(`epochs × rollout_len`).  A larger rollout gives the GAE a longer
horizon before bootstrap and more data per update, but also means more
gradient steps per rollout: at 50000 you'd run `4 × 50000/64 = 3125`
updates on one batch, which pushes PPO toward overfitting a single
rollout.

### PPO flag reference

| Flag | Default | Meaning |
|---|---|---|
| `--ppo_gamma` | 0.99 | Discount factor for GAE / returns |
| `--ppo_lam` | 0.95 | GAE trace-decay (bias-variance tradeoff) |
| `--ppo_clip_eps` | 0.2 | Clipping epsilon of the clipped surrogate |
| `--ppo_epochs` | 4 | SGD passes over each rollout |
| `--ppo_mini_batch_size` | 64 | Mini-batch size for PPO updates |
| `--ppo_entropy_coef` | 0.01 | Entropy bonus coefficient (see SB3-tricks section) |
| `--ppo_value_coef` | 0.5 | Value-loss coefficient |
| `--ppo_vf_clip_eps` | None | Value-function clipping (None = unclipped) |
| `--ppo_target_kl` | None | Early-stop a PPO epoch when approx-KL exceeds this (see SB3-tricks section) |
| `--ppo_rollout_len` | 2048 | Env steps collected per training epoch |
| `--ppo_discrete` / `--ppo_continuous` | discrete | Action-space selection |
| `--ppo_log_std_init` | log(0.5) ≈ −0.693 | Initial log_std of the continuous Gaussian policy (σ₀ = 0.5) |
| `--ppo_log_std_min` | −10.0 | Lower safety clamp of log_std (σ floor ≈ 4.5e-5).  Kept far below a converged policy's operating range so the entropy gradient never dies at the boundary (the old −2.0 default froze exploration at σ = 0.135) |
| `--ppo_log_std_max` | 2.0 | Upper safety clamp of log_std (σ ceiling ≈ 7.4; bounds entropy) |

The `log_std_*` knobs control the **continuous** Gaussian policy's
standard deviation.  `log_std_init` sets where exploration starts;
`log_std_min`/`log_std_max` are wide safety rails (SOTA pattern: SB3
unbounded, rsl_rl `std_range=(1e-6, 1e6)`, jaxrl `LOG_STD_MIN=-10`) that
prevent σ collapse/explosion while leaving entropy freely trainable.

### Stable-Baselines3-style tricks

The PPO implementation supports three opt-in techniques that
Stable-Baselines3 bakes into its default algorithms.  All three are
**off by default** — every existing flag combination behaves exactly as
before — and each is documented with *when it helps* and *when it
doesn't*.

#### Reward normalization (`--reward_normalize`)

Normalizes rewards by the running standard deviation of discounted
returns, matching SB3 `VecNormalize(norm_reward=True)`.  For each step,
`running_return = gamma * running_return + reward` is fed to a Welford
`RunningMeanStd`; the reward is then divided by `sqrt(var + eps)` and
clamped to `[-reward_norm_clip, reward_norm_clip]` (default 10.0).

- **When it helps:** environments with large or unnormalized reward
  scales (Pendulum-v1, MuJoCo) where an O(1000) value loss otherwise
  dominates the total loss and swamps the entropy/policy signal.
- **When it doesn't:** rewards already O(1) (CartPole, most toy envs) —
  normalization is harmless but unnecessary.
- **Eval always uses raw rewards** (SB3 `VecNormalize.train(False)`):
  the running-return stats are not updated during evaluation, so
  `eval_return` stays comparable across runs.  The running stats travel
  with the checkpoint and are restored on resume.

#### Orthogonal initialization (`--param_init ortho`)

Applies SB3-style orthogonal initialization to the model after
construction: gain `sqrt(2)` for policy/hidden layers, `1.0` for
value/head layers, zero biases — matching the original PPO
implementation.  `log_std` is left untouched.  Works for any method
(PPO, SAC, and non-RL MLPs alike).

- **When it helps:** continuous RL policies, where the default PyTorch
  init can leave the policy near-degenerate at start and slow the
  critical early phase.
- **When it doesn't:** fine-tuned / checkpoint-loaded models (init is
  skipped when weights are loaded from a checkpoint) or discrete
  policies that already train fine with defaults.

#### `target_kl` early stopping (`--ppo_target_kl <float>`)

SB3 stops a PPO epoch early when the approximate KL divergence (the
Schulman **k1** estimator, `0.5 * mean(log_ratio^2)`, always >= 0)
exceeds `target_kl`.  This guards against destructive updates on
high-variance minibatches.

- **Suggested starting point:** `--ppo_target_kl 0.03`.
- **Too low** (e.g. < 0.005) stalls learning; **too high** (e.g. > 0.1)
  disables the safety.  `None` (default) disables early stopping.
- Only affects PPO epochs after the first rollout, so it is a pure
  safety valve for the on-policy update phase.

#### Entropy coefficient (`--ppo_entropy_coef <float>`)

Coefficient of the entropy bonus (default 0.01).  SB3 uses **0.0** for
continuous PPO, relying on init + normalization instead; genml_kit keeps
a small positive default as light regularization.

- **Higher** (0.05–0.1) encourages broader exploration — may help hard
  exploration tasks.
- **Lower** (toward 0.0) sharpens the policy — may help fine-tuning late
  in training.
- No single value is best for all tasks: sweep 0.0–0.1 on your task.

A typical "serious" continuous-PPO setup (Pendulum-style) looks like:

```
genml-kit-train --pipeline rl --method ppo --env_id Pendulum-v1 \
    --ppo_continuous --obs_normalize --reward_normalize --param_init ortho \
    --ppo_target_kl 0.03 --epochs 500
```

### Custom environments

Pass `--env_script /path/to/env.py` to load a custom environment.  The
script must expose a `make_env()` factory returning a Gymnasium-style env
(`reset()`, `step(action)`, `close()`; `observation_space` and
`action_space` attributes, which the pipeline introspects to size the
networks).  `make_env()` may take a single `seed` keyword argument if your
environment wants reproducible initialization.

### Environment & replay options

Relevant flags:
  - `--env_id`, `--obs_dim`, `--env_script` -- environment selection.
  - `--replay_capacity`, `--batch_size`, `--n_step`, `--prioritized`,
    `--per_alpha`, `--per_beta_start`, `--per_beta_frames` -- replay tuning.
  - `--obs_normalize`, `--obs_norm_clip` -- running observation
    normalization (Welford).  Both train and eval observations are
    normalized with the same RMS statistics.
  - `--reward_normalize`, `--reward_norm_clip` -- running-return reward
    normalization (SB3 `VecNormalize(norm_reward=True)`).  See
    "Stable-Baselines3-style tricks" below.
  - `--param_init ortho` -- SB3-style orthogonal parameter initialization.
    See "Stable-Baselines3-style tricks" below.
  - `--ppo_gamma`, `--ppo_lam`, `--ppo_clip_eps`, `--ppo_epochs`,
    `--ppo_mini_batch_size`, `--ppo_value_coef`, `--ppo_vf_clip_eps`,
    `--ppo_rollout_len`, `--ppo_discrete`/`--ppo_continuous`,
    `--ppo_log_std_init/min/max` -- PPO update and policy knobs.  See the
    "PPO flag reference" table above.
  - `--ppo_target_kl` -- early-stop a PPO epoch when the k1 approximate-KL
    exceeds this threshold (SB3 `target_kl`).  See below.
  - `--ppo_entropy_coef` -- entropy bonus coefficient (default 0.01).
    See below.
  - `--eval_episodes`, `--env_seed` -- evaluation.
  - `--record_eval_video` -- record one MP4 (or GIF fallback) per evaluation
    episode under `<checkpoint_dir>/videos/`; only when the environment
    supports rendering (see "Evaluation videos" below).

### Episode-end semantics (terminated vs truncated)

Gymnasium distinguishes *terminated* (the MDP truly ended, e.g. the pole
fell) from *truncated* (the episode was cut short, e.g. by a time limit).
genml_kit threads the true `terminated` flag through the replay buffer,
rollout buffer, GAE, and the TD targets: a truncated final step still
bootstraps `γ V(s′)`, while a terminated one does not.  Environments using
the legacy 4-tuple gym API (no `terminated`/`truncated` split) fall back to
treating `done` as `terminated`.

## References And Further Reading

- Xie et al., [SimMIM: A Simple Framework for Masked Image Modeling](https://arxiv.org/abs/2111.09886), CVPR 2022.
- Assran et al., [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://arxiv.org/abs/2301.08243), CVPR 2023.
- Grill et al., [Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning](https://arxiv.org/abs/2006.07733), NeurIPS 2020.
  The EMA teacher and stop-gradient design used by I-JEPA-style methods.
- Caron et al., [Emerging Properties in Self-Supervised Vision Transformers](https://arxiv.org/abs/2104.14294), ICCV 2021.
  The DINO paper — self-distillation with an EMA teacher and multi-crop.
  The centering + sharpening trick to avoid representation collapse.
- Khosla et al., [Supervised Contrastive Learning](https://arxiv.org/abs/2004.11362), NeurIPS 2020.
- He et al., [Masked Autoencoders Are Scalable Vision Learners](https://arxiv.org/abs/2111.06377), CVPR 2022.
  A useful comparison for masked-image pre-training.
- Hu et al., [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685), ICLR 2022.
- Howard & Ruder, [Universal Language Model Fine-tuning for Text Classification](https://arxiv.org/abs/1801.06146), ACL 2018.
  Introduced layer-wise learning rate decay.
- Zhang et al., [mixup: Beyond Empirical Risk Minimization](https://arxiv.org/abs/1710.09412), ICLR 2018.
- Lin et al., [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002), ICCV 2017.
  Introduced focal loss and the $`(1-p_t)^\gamma`$ easy-example down-weighting.
- Chen & Guestrin, [XGBoost: A Scalable Tree Boosting System](https://arxiv.org/abs/1603.02754), KDD 2016.
- The [PyTorch optimization documentation](https://pytorch.org/docs/stable/optim.html)
  explains AdamW, schedulers, and gradient clipping.
- The [scikit-learn metrics documentation](https://scikit-learn.org/stable/modules/model_evaluation.html)
  is useful when choosing metrics for imbalanced classification.

## Development

```bash
git clone https://github.com/davidel/genml_kit
cd genml_kit
pip install -e ".[dev,all]"
pytest
```

Tests run with warnings-as-errors (`filterwarnings = ["error"]` in
`pyproject.toml`), so a new dependency deprecation warning fails the suite
instead of scrolling past.  Code is formatted with
[yapf](https://github.com/google/yapf) using the project's `.style.yapf`
(Google style, 2-space indent) and linted with
[Ruff](https://docs.astral.sh/ruff/) (`ruff check .`).

### Naming Conventions For Class Members

- Anything that PyTorch registers — child `nn.Module`, `nn.Parameter`,
  registered buffer — is **never** underscore-prefixed, regardless of
  intended visibility.  A leading underscore would change `state_dict()`
  keys, optimizer param groups, and DDP behavior.
- Everything else that is not part of a class's public contract
  (internal state, helper hooks invoked only by the class itself) gets a
  leading underscore.  A member called from other classes, overridden as
  an extension point, or documented as API stays public.
- Extension-point methods that *callers* invoke (e.g.
  `PretrainMethod.build_transform`) are public; same-named hooks the
  base class *itself* dispatches (e.g. `BaseImageProcessor._build_transform`)
  are private.

## License

Apache-2.0

