nidl.estimators: Available estimators

This modules details the public API you should use and implement for a nidl compatible estimator, as well as the estimators available in nidl.

Introduction

An estimator is an object that fits a model based on some training data and is capable of inferring some properties on new data. It can be, a classifier, a clustering algorithm, a regressor or a transformer (in the “scikit-learn” sense). All estimators implement a fit method. Behind the hood, it inherits from a pytorch_lightning.LightningModule, and thus benefits from all the features of the pytorch_lightning library:

  • Distributed multi-GPU training

  • Logging and visualization

  • Clear organization of the training and evaluation code

  • Automatic checkpointing and early stopping

  • Callback logic

Instanciation

The estimator __init__ method might accept constants as arguments that determine the estimator’s behavior (like the hyperparameters and training settings). It should not, however, take the actual training data as an argument, as this is left to the fit method.

Fitting

The next thing you will probably want to do is to estimate some parameters in the model. This is implemented in the fit method, and it’s where the training happens.

The fit method takes the following training data as arguments:

Parameters

train_dataloader

torch DataLoader [(n_samples, *)]

val_dataloader

torch DataLoader [(n_samples, *)]

Build as a LightningModule, the fit method gets organized under a training_step and validation_step methods.

Estimator types

The proposed types of estimators are transformers, classifiers, regressors, and clustering algorithms.

Transformers inherit from TransformerMixin, and implement a transform method. These are estimators which take the input, and transform it in some way. Note that they should never change the number of input samples, and the output of transform should correspond to its input samples in the same given order.

Regressors inherit from RegressorMixin, and implement a predict method returning the values assigned to newly given samples. In this case the training data must returns two tensors.

Classifiers inherit from ClassifierMixin, and implement a predict method returning the labels assigned to newly given samples. In this case the training data must returns two tensors.

Clustering inherit from ClusterMixin, and implement a predict method returning the labels assigned to newly given samples. In this case the training data must returns two tensors.

Build as a LightningModule, the transform and the predict method gets organized under a the transform_step and the predict_step methods.

Base Classes

Base classes for all nidl estimators.

BaseEstimator([callbacks, ...])

Base class for all estimators in the NIDL framework designed for scalability.

ClassifierMixin()

Mixin class for all classifiers in nidl.

ClusterMixin()

Mixin class for all cluster estimators in nidl.

RegressorMixin()

Mixin class for all regression estimators in nidl.

TransformerMixin()

Mixin class for all transformers in nidl.

 
        classDiagram
  LightningModule <|-- BaseEstimator

Self-Supervised Learning

Self-supervised learning embedding estimators.

Embedding estimators

SimCLR(encoder[, encoder_kwargs, ...])

SimCLR [Rfa9275c25f8b-1].

DCL(encoder[, encoder_kwargs, ...])

Decoupled Contrastive Learning [R23fac62c525a-1].

YAwareContrastiveLearning(encoder[, ...])

y-Aware Contrastive Learning [R8ad1d13554f2-1].

BarlowTwins(encoder[, encoder_kwargs, ...])

Barlow Twins [Rf18316f97edd-1].

DINO(encoder[, encoder_kwargs, ...])

DINO [Rf531b9701544-1].

IJEPA(encoder[, dim, context_block_scale, ...])

Implementation of I-JEPA [Re5620117445e-1].

NeuroJEPA(encoder[, mask_scale_configs, ...])

Implementation of Neuro-JEPA [R523d24230123-1].

 
        classDiagram
  BaseEstimator <|-- BarlowTwins
  BaseEstimator <|-- DCL
  BaseEstimator <|-- DINO
  BaseEstimator <|-- IJEPA
  BaseEstimator <|-- NeuroJEPA
  BaseEstimator <|-- SimCLR
  BaseEstimator <|-- YAwareContrastiveLearning
  TransformerMixin <|-- BarlowTwins
  TransformerMixin <|-- DCL
  TransformerMixin <|-- DINO
  TransformerMixin <|-- IJEPA
  TransformerMixin <|-- NeuroJEPA
  TransformerMixin <|-- SimCLR
  TransformerMixin <|-- YAwareContrastiveLearning

See nidl.losses: Available losses and nidl.estimators.ssl.utils: Available projection heads for the losses and projection heads used by these estimators.

Probing estimators

Probing estimators, which are used to evaluate the quality of the representations learned by an embedding estimator.

ModelProbing(embedding_estimator, probe[, ...])

Estimator to probe the representation of an embedding estimator.

Autoencoders

Autoencoder estimators and losses.

Embedding estimators

VAE(encoder, decoder, encoder_out_dim, ...)

Variational Auto-Encoder (VAE) [Rfc808ee91d8e-1] [Rfc808ee91d8e-2].

 
        classDiagram
  BaseEstimator <|-- VAE
  TransformerMixin <|-- VAE

See nidl.losses: Available losses for the losses used by these estimators.

Dummy estimator

A dummy estimator that does not perform any training. This can be useful for testing purposes.

DummyEmbeddingEstimator([strategy, ...])

A dummy embedding estimator returning an embedding independent of the input data.