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 estimators, losses and associated tools.

Estimators

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

SimCLR [Rfa9275c25f8b-1].

YAwareContrastiveLearning(encoder[, ...])

y-Aware Contrastive Learning [R8ad1d13554f2-1].

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

Barlow Twins [Rf18316f97edd-1].

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

DINO [Rf531b9701544-1].

 
        classDiagram
  BaseEstimator <|-- BarlowTwins
  BaseEstimator <|-- DINO
  BaseEstimator <|-- SimCLR
  BaseEstimator <|-- YAwareContrastiveLearning
  TransformerMixin <|-- BarlowTwins
  TransformerMixin <|-- DINO
  TransformerMixin <|-- SimCLR
  TransformerMixin <|-- YAwareContrastiveLearning

Losses

InfoNCE([temperature])

Implementation of the InfoNCE loss [Re38fc64e0ed4-1], [Re38fc64e0ed4-2].

YAwareInfoNCE([kernel, bandwidth, temperature])

Implementation of the y-Aware InfoNCE loss [Ra2feb9ab43ec-1].

BarlowTwinsLoss([lambd])

Implementation of the Barlow Twins loss [Re83c9b545e4a-1].

DINOLoss([output_dim, warmup_teacher_temp, ...])

Implementation of the DINO loss [Re2e7efabd714-1].

Projection heads

ProjectionHead(blocks)

Base class for all projection and prediction heads in self-supervised estimators.

SimCLRProjectionHead([input_dim, ...])

Projection head used for SimCLR.

YAwareProjectionHead([input_dim, ...])

Projection head used for yAware contrastive learning.

BarlowTwinsProjectionHead([input_dim, ...])

Projection head used for Barlow Twins [R86bff2d09119-1].

DINOProjectionHead([input_dim, hidden_dim, ...])

Projection head used in DINO [Re0843dd52443-1].

Autoencoders

Autoencoder estimators and losses.

Estimators

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

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

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

Losses

BetaVAELoss([beta, default_dist])

Compute the Beta-VAE loss [Rd208eb1f92d3-1].