Reference training

train.trainer

EpochCallback

Bases: Protocol

Interface for callbacks called by the trainer after each epoch.

BasicTrainer(dataset, model, objective, optimizer, epochs, scheduler=None, device='cuda', batch_size=128, num_workers=1, accumulation_steps=1, epoch_callback=None)

Implements basic training loop for Pytorch models. Checkpoints includes random states - any restarts from checkpoint preservers reproducibility.

Parameters:

Name	Type	Description	Default
dataset	ImageDataset	Training dataset that gives (x, y) tensor pairs.	required
model	Module	Pytorch nn.Module for model / backbone.	required
objective	Module	Pytorch nn.Module for objective / loss function.	required
optimizer	Optimizer	Pytorch optimizer.	required
epochs	int	Number of training epochs.	required
scheduler	(_LRScheduler, optinal)	Pytorch scheduler.	None
device	str	Device to be used for training.	'cuda'
batch_size	int	Training batch size.	128
num_workers	int	Number of data loading workers in torch DataLoader.	1
accumulation_steps	int	Number of gradient accumulation steps.	1
epoch_callback	EpochCallback | None	Callback function to be called after each epoch.	None

get_random_states()

Gives dictionary of random states for reproducibility.

train.objective

GSAPLoss(temperature=0.001, negative_drop_rate=0.9, chunk_size=64)

Bases: SmoothAPLoss

Global sigmoid AP surrogate (gSAP), arXiv:2610.09863.

All valid directed pairs are ranked together, including negatives from queries without positives. Optional negative dropping removes the lowest negative scores and preserves all positives. Defaults follow the paper's distractor-query recipe; set negative_drop_rate=0 for the full loss.

Parameters:

Name	Type	Description	Default
temperature	float	Positive sigmoid temperature, default 0.001.	0.001
negative_drop_rate	float	Fraction of lowest negatives to drop, in [0, 1].	0.9
chunk_size	int	Number of positive ranks computed per checkpointed chunk.	64

Configure global ranking temperature, negative dropping, and chunk size.

from_similarity(similarities, positive_mask, valid_mask=None)

Optimize AP over the flattened set of valid pairs.

SigLIPLoss(scale=10.0, bias=-10.0, learnable=True)

Bases: Module

Sigmoid pairwise objective with a shared learnable scale and bias.

Uses cosine similarities and excludes self-pairs for embedding batches. from_similarity also supports rectangular/cross-modal matrices. The signed logistic loss is summed over valid pairs and divided by the number of queries, as in SigLIP; it is not averaged over all pairs. Negative-only batches still train. No-positive batches behave differently from AP losses.

Initialize the logit scale and bias, optionally making them trainable.

forward(embeddings, labels)

Compute sigmoid pair loss from labeled embeddings, excluding self-pairs.

from_similarity(similarities, positive_mask, valid_mask=None)

Sum signed logistic losses over valid pairs and divide by query count.

SmoothAPLoss(temperature=0.01, chunk_size=64)

Bases: Module

Per-query sigmoid AP surrogate (mSAP).

forward(embeddings, labels) uses cosine similarities and excludes self-pairs. from_similarity(S, Y, valid_mask=None) accepts rectangular similarity/relevance matrices, including cross-modal diagonal positives. Queries without valid positives are omitted from the mean. Empty batches return a differentiable zero. Each rank includes sigmoid(0) = 0.5.

Parameters:

Name	Type	Description	Default
temperature	float	Positive sigmoid temperature.	0.01
chunk_size	int	Number of positive ranks computed per checkpointed chunk.	64

Set the sigmoid temperature and positive-comparison chunk size.

forward(embeddings, labels)

Compute the AP surrogate from labeled embeddings, excluding self-pairs.

from_similarity(similarities, positive_mask, valid_mask=None)

Optimize the mean AP of rows with at least one valid positive.

ArcFaceLoss(num_classes, embedding_size, margin=0.5, scale=64)

Bases: Module

Wraps Pytorch Metric Learning ArcFaceLoss.

Parameters:

Name	Type	Description	Default
num_classes	int	Number of classes.	required
embedding_size	int	Size of the input embeddings.	required
margin	int	Margin for ArcFace loss (in radians).	0.5
scale	int	Scale parameter for ArcFace loss.	64

SubCenterArcFaceLoss(num_classes, embedding_size, margin=0.5, scale=64, sub_centers=3)

Bases: Module

ArcFace with multiple centers per class; margin is in radians.

Parameters:

Name	Type	Description	Default
num_classes	int	Number of training classes.	required
embedding_size	int	Embedding dimension.	required
margin	float	Angular margin in radians, matching the ArcFace wrapper.	0.5
scale	float	Logit scale.	64
sub_centers	int	Number of centers per class.	3

TripletLoss(margin=0.2, mining='semihard', distance='l2_squared')

Bases: Module

Wraps Pytorch Metric Learning TripletMarginLoss.

Parameters:

Name	Type	Description	Default
margin	int	Margin for triplet loss.	0.2
mining	str	Type of triplet mining. One of: 'all', 'hard', 'semihard'	'semihard'
distance	str	Distance metric for triplet loss. One of: 'cosine', 'l2', 'l2_squared'	'l2_squared'

SoftmaxLoss(num_classes, embedding_size)

Bases: Module

CE with single dense layer classification head.

Parameters:

Name	Type	Description	Default
num_classes	int	Number of classes.	required
embedding_size	int	Size of the input embeddings.	required

MinedLoss(loss, miner)

Bases: Module

Apply a PyTorch Metric Learning miner before a compatible backend loss.

Both modules are registered so proxies and any trainable loss parameters remain available to the optimizer and checkpoint state.

CombinedLoss(objectives, weights=None)

Bases: Module

Weighted sum of named objective(embeddings, labels) modules.

Weights are stored in checkpoint state. Include this module's parameters in the optimizer to train constituent classifiers, proxies, or scales.

train.ranking

Differentiable per-query and global average precision objectives.

The sigmoid ranks follow Eqs. (3), (8)--(10) of Psomas et al., "Global Average Precision for Representation Learning", arXiv:2610.09863.

SmoothAPLoss(temperature=0.01, chunk_size=64)

Bases: Module

Per-query sigmoid AP surrogate (mSAP).

forward(embeddings, labels) uses cosine similarities and excludes self-pairs. from_similarity(S, Y, valid_mask=None) accepts rectangular similarity/relevance matrices, including cross-modal diagonal positives. Queries without valid positives are omitted from the mean. Empty batches return a differentiable zero. Each rank includes sigmoid(0) = 0.5.

Parameters:

Name	Type	Description	Default
temperature	float	Positive sigmoid temperature.	0.01
chunk_size	int	Number of positive ranks computed per checkpointed chunk.	64

Set the sigmoid temperature and positive-comparison chunk size.

forward(embeddings, labels)

Compute the AP surrogate from labeled embeddings, excluding self-pairs.

from_similarity(similarities, positive_mask, valid_mask=None)

Optimize the mean AP of rows with at least one valid positive.

GSAPLoss(temperature=0.001, negative_drop_rate=0.9, chunk_size=64)

Bases: SmoothAPLoss

Global sigmoid AP surrogate (gSAP), arXiv:2610.09863.

All valid directed pairs are ranked together, including negatives from queries without positives. Optional negative dropping removes the lowest negative scores and preserves all positives. Defaults follow the paper's distractor-query recipe; set negative_drop_rate=0 for the full loss.

Parameters:

Name	Type	Description	Default
temperature	float	Positive sigmoid temperature, default 0.001.	0.001
negative_drop_rate	float	Fraction of lowest negatives to drop, in [0, 1].	0.9
chunk_size	int	Number of positive ranks computed per checkpointed chunk.	64

Configure global ranking temperature, negative dropping, and chunk size.

from_similarity(similarities, positive_mask, valid_mask=None)

Optimize AP over the flattened set of valid pairs.

SigLIPLoss(scale=10.0, bias=-10.0, learnable=True)

Bases: Module

Sigmoid pairwise objective with a shared learnable scale and bias.

Uses cosine similarities and excludes self-pairs for embedding batches. from_similarity also supports rectangular/cross-modal matrices. The signed logistic loss is summed over valid pairs and divided by the number of queries, as in SigLIP; it is not averaged over all pairs. Negative-only batches still train. No-positive batches behave differently from AP losses.

Initialize the logit scale and bias, optionally making them trainable.

forward(embeddings, labels)

Compute sigmoid pair loss from labeled embeddings, excluding self-pairs.

from_similarity(similarities, positive_mask, valid_mask=None)

Sum signed logistic losses over valid pairs and divide by query count.