idtrackerai.base.tracker.contrastive#

Warning

The code reference is a work in progress and may contain inconsistencies.

Contrastive Learning module

class idtrackerai.base.tracker.contrastive.PairsOfFragments(*args: Any, **kwargs: Any)[source]#

Dataset with all pairs of fragments (positive and negative) which returns only the locations of selected images (the images are loaded in collate_fun).

pairs#

List of fragment pairs.

Type:: list[tuple[Fragment, Fragment]]

__init__(pairs: list[tuple[Fragment, Fragment]]) → None[source]#

__len__() → int[source]#: Return the number of pairs in the dataset.

__getitem__(pair_index: int) → tuple[tuple[int, int], tuple[int, int], int][source]#

Get two images from a specific pair of fragments.

Parameters:

pair_index (int) – Index of the pair to retrieve.

Returns:

A tuple containing:

Locations of an image from the first fragment.
Locations of an image from the second fragment.
Index of the pair.

Return type:

tuple

class idtrackerai.base.tracker.contrastive.BatchSampler(*args: Any, **kwargs: Any)[source]#

Custom implementation of a torch.utils.data.BatchSampler.

This sampler generates batches of indices based on a probability distribution derived from loss scores and pair sizes. The probabilities can be updated dynamically during training.

batch_size#

Number of samples in each batch.

Type:: int

n_batches#

Number of batches to generate.

Type:: int | None

negative_pairs_sizes#

Sizes of negative pairs.

Type:: Tensor

positive_pairs_sizes#

Sizes of positive pairs.

Type:: Tensor

negative_probabilities#

Sampling probabilities for negative pairs.

Type:: Tensor

positive_probabilities#

Sampling probabilities for positive pairs.

Type:: Tensor

__init__(negative_pairs_sizes: object, positive_pairs_sizes: object, negative_loss_scores: object, positive_loss_scores: object, batch_size: int, n_batches: int | None = None) → None[source]#

__iter__() → Iterator[list[int]][source]#

Generate batches of indices.

Yields:: list[int] – A batch of indices.

update_probabilities(negative_scores: object, positive_scores: object) → None[source]#

Update the sampling probabilities for negative and positive pairs based on their loss scores. The parameter lpha in the 2025 article is hardcoded to 0.5 so both probability contributions are summed with the same weight.

Parameters:

negative_scores (Tensor) – Loss scores for negative pairs.
positive_scores (Tensor) – Loss scores for positive pairs.

class idtrackerai.base.tracker.contrastive.ContrastiveDataLoader(*args, **kwargs)[source]#

Protocol class for better typing our DataLoaders

__init__(*args, **kwargs)[source]#

classmethod __subclasshook__(other)[source]#

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

class idtrackerai.base.tracker.contrastive.ContrastiveLearning(fragments: ListOfFragments, saving_folder: Path | None = None, first_gfrag: GlobalFragment | None = None, batch_size: int = 400, preload_images_max_mbytes: float | None = None, learning_rate: float = 0.001, embedding_dimensions: int = 8, min_fragment_length: int = 4)[source]#

Main API class for Contrastive Learning.

model: ResNet18#: RasNet18 model

optimizer: torch.optim.optimizer.Optimizer#: Optimizer

val_loader: ContrastiveDataLoader#: Validation DataLoader

train_loader: ContrastiveDataLoader#: Contrastive training DataLoader

gfrag_loader: ContrastiveDataLoader | None#: DataLoader for a single Global Fragments images used to initialize kmeans clusters. None if there are no Global Fragments in the video.

__init__(fragments: ListOfFragments, saving_folder: Path | None = None, first_gfrag: GlobalFragment | None = None, batch_size: int = 400, preload_images_max_mbytes: float | None = None, learning_rate: float = 0.001, embedding_dimensions: int = 8, min_fragment_length: int = 4) → None[source]#

saving_folder: Path#: Saving folder for checkpoints

learning_rate: float#: Optimizer learning rate

embedding_dimensions: int#: Number of dimensions of the embedded space

n_animals: int#: Number of animals in the video

batch_size: int#: Number of pairs of each kind of images (positive and negative) used in a single training batch

n_negative_pairs: int#: The number of negative pairs of Fragments we have

loss_scores: object#: Sequence of floats representing the loss scores of every pair of Fragments used in contrastive. Loss scores increase when a pair of images is sampled from a specific pair of Fragments and its loss is non zero.

static criterion(embedded_A: object, embedded_B: object, first_positive: int, margin: float = 10) → object[source]#: Pairwise distance loss criterion. Negative pairs are pushed away until they are at distance margin. Positive pairs are pulled together until they are at distance 1

set_model(weights_path: Path | None = None, compile: bool = False) → None[source]#: Initializes the contrastive model from a knowledge transfer file or from scratch

train_step(n_batches: int, output: ~collections.abc.Callable[[str], None] = <function ContrastiveLearning.<lambda>>, starting_batch_number: int = 0) → tuple[float, float][source]#

Perform n_batches training steps.

Parameters:

n_batches (int) – Number of batches to train on.
output (Callable[[str], None]) – Function to output training progress.
starting_batch_number (int) – Starting batch number for logging.

Returns:

Fraction of positive and negative sampled pairs with loss.

Return type:

tuple[float, float]

train(check_every: int | None = None, skipped_validations: int = 5, patience: int = 30, target_silhouette_score: float = 0.91) → float[source]#: Main method to train the contrastive

idtrackerai.base.tracker.contrastive.collate_fun(batch: list[tuple[tuple[int, int], tuple[int, int], int]], images_sources: Sequence[Dataset | Path | str | ndarray | H5DatasetProxy]) → list[object][source]#: Receives the batch with N groups of images locations (episode and index) and a label. These are used to load the images and generate the batch tensor

idtrackerai.base.tracker.contrastive.silhouette_scores(X: object, labels: object) → object[source]#

Silhouette score implemented in PyTorch for GPU acceleration