embedding_benchmark

Attributes

args
parser

Classes

Baseline	The main class that runs the baseline methods.
BaselineMethod	An abstract class that holds common code of our baseline methods.
Config
KNNClassifier	A lightly KNN Classifier modified to log mean average precision metric.
LinearClassifier	A lightly Linear Classifier, modified to log the mean average precision
MetricCallback	A [Lightly] Callback that collects log metrics from the LightningModule and stores them after
MetricModule
SwinL384	A lightly Linear Classifier, modified to log the mean average precision
SwinL384Baseline	An abstract class that holds common code of our baseline methods.
ViTEmbedding	This module is used to extract features from the Vision Transformer Classifier in eval mode
ViT_B_16Baseline	An abstract class that holds common code of our baseline methods.
ViT_B_16Classifier	A fully supervised model that uses the Vision Transformer model from the torchvision library

Functions

clear_cache()
knn_predict(→ torch.Tensor)	[Modified version from lightly, which returns the scores instead of the predictions]
timing_decorator(func)

Module Contents

class embedding_benchmark.Baseline

The main class that runs the baseline methods.

aggregate_metrics(cfg: Config)

calculate_mean_std(df: pandas.DataFrame, groupby: List[str])

Calculate the mean and standard deviation for all numerical columns grouped by two identifiers, and return a DataFrame with the results in the ± notation.

Parameters:

• df (pd.DataFrame) – Input DataFrame.

• groupby (List[str]) – List of which columns to group by.

Returns:

DataFrame with mean and standard deviation in ± notation.

Return type:

pd.DataFrame

run(args)

Run the class as specified in the config.

args: argparse.Namespace - The CLI arguments, used as kwargs to instantiate a Config object.

methods: Dict[str, Type[BaselineMethod]]

class embedding_benchmark.BaselineMethod(cfg: Config)

An abstract class that holds common code of our baseline methods.

The class runs:

• embedding training

• kNN evaluation

• linear evaluation

Reported metrics are:

• Top-1 accuracy

• Top-5 accuracy

• Mean Average Precision (mAP)

The baseline method can be configured by inheriting from this class and overriding specific attributes or functions as well as passing a config object.

embedding_training()

get_embedding_model() → torch.nn.Module

Must return a model that returns features on forward pass.

knn_eval() → None

Runs KNN evaluation on the given model.

Parameters follow InstDisc [0] settings.

The most important settings are:

• Num nearest neighbors: 200

• Temperature: 0.1

References: - [0]: InstDict, 2018, https://arxiv.org/abs/1805.01978

linear_eval() → None

Runs a linear evaluation on the given model.

Parameters follow SimCLR [0] settings.

The most important settings are:

• Backbone: Frozen

• Epochs: 90

• Optimizer: SGD

• Base Learning Rate: 0.1

• Momentum: 0.9

• Weight Decay: 0.0

• LR Schedule: Cosine without warmup

References

• [0]: SimCLR, 2020, https://arxiv.org/abs/2002.05709

run_baseline_method()

train(classifier, epochs, train_dataset, val_dataset, log_name)

cfg: Config

embedding_train_dataset: Iterable[Tuple[torch.Tensor, torch.Tensor]]

embedding_train_transform

embedding_val_dataset: Iterable[Tuple[torch.Tensor, torch.Tensor]]

eval_train_transform

feature_dim: int = 2048

knn_train_dataset: Iterable[Tuple[torch.Tensor, torch.Tensor]]

knn_val_dataset: Iterable[Tuple[torch.Tensor, torch.Tensor]]

linear_train_dataset: Iterable[Tuple[torch.Tensor, torch.Tensor]]

linear_val_dataset: Iterable[Tuple[torch.Tensor, torch.Tensor]]

method_dir

method_specific_augmentation

model: torch.nn.Module

property name: str

normalize_transform

resize_transform

skip_embedding_training: bool = False

val_transform

class embedding_benchmark.Config

accelerator: str = 'auto'

aggregate_metrics: bool = True

baseline_id: str | None = None

batch_size_per_device: int = 16

check_val_every_n_epoch: int = 5

checkpoint_path: pathlib.Path | None = None

devices: int = 1

epochs: int = 200

experiment_result_metrics: List[str] | None = []

log_dir: pathlib.Path

methods: List[str] | None = None

num_classes: int = 50

num_workers: int = 4

precision: str = '16-mixed'

profile = None

skip_embedding_training: bool = False

skip_knn_eval: bool = False

skip_linear_eval: bool = False

test_run: bool = True

class embedding_benchmark.KNNClassifier(model: torch.nn.Module, num_classes: int, knn_k: int = 200, knn_t: float = 0.1, feature_dtype: torch.dtype = torch.float32, normalize: bool = True)

Bases: MetricModule

A lightly KNN Classifier modified to log mean average precision metric. Also it now inherits from MetricModule and the logging logic has changed.

configure_optimizers() → None

on_train_epoch_start() → None

on_validation_end() → None

on_validation_epoch_start() → None

training_step(batch, batch_idx) → None

validation_step(batch, batch_idx) → None

feature_dtype

knn_k = 200

knn_t = 0.1

model

normalize = True

num_classes

class embedding_benchmark.LinearClassifier(model: torch.nn.Module, batch_size_per_device: int, feature_dim: int, num_classes: int, freeze_model: bool = False, enable_logging: bool = True)

Bases: MetricModule

A lightly Linear Classifier, modified to log the mean average precision Also, the logging logic has changed + it now inherits from MetricModule Further, the LinearClassifier now also allows the instantiation of fully supervised models.

build_classification_head(feature_dim: int, num_classes: int)

build_critierion()

configure_optimizers() → Tuple[List[torch.optim.Optimizer], List[Dict[str, Any | str]]]

forward(images: torch.Tensor) → torch.Tensor

on_train_epoch_start() → None

training_step(batch: Tuple[torch.Tensor, Ellipsis], batch_idx: int) → torch.Tensor

validation_step(batch: Tuple[torch.Tensor, Ellipsis], batch_idx: int) → torch.Tensor

batch_size_per_device

classification_head

criterion

enable_logging = True

feature_dim

freeze_model = False

model

num_classes

class embedding_benchmark.MetricCallback

Bases: pytorch_lightning.callbacks.Callback

A [Lightly] Callback that collects log metrics from the LightningModule and stores them after every epoch.

train_metrics

Dictionary that stores the last logged metrics after every train epoch.

val_metrics

Dictionary that stores the last logged metrics after every validation epoch.

on_train_end(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule) → None

on_validation_end(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule) → None

train_metrics: Dict[str, List[float]]

val_metrics: Dict[str, List[float]]

class embedding_benchmark.MetricModule(num_classes: int)

Bases: pytorch_lightning.LightningModule

on_train_epoch_end()

on_validation_epoch_end()

update_train_metrics(pred_scores: torch.Tensor, targets: torch.Tensor)

update_val_metrics(pred_scores: torch.Tensor, targets: torch.Tensor)

enable_logging = True

num_classes

class embedding_benchmark.SwinL384(batch_size_per_device, feature_dim, num_classes)

Bases: LinearClassifier

A lightly Linear Classifier, modified to log the mean average precision Also, the logging logic has changed + it now inherits from MetricModule Further, the LinearClassifier now also allows the instantiation of fully supervised models.

build_critierion()

configure_optimizers()

forward(x: torch.Tensor) → torch.Tensor

enable_logging: bool = False

class embedding_benchmark.SwinL384Baseline(args)

Bases: BaselineMethod

An abstract class that holds common code of our baseline methods.

The class runs:

• embedding training

• kNN evaluation

• linear evaluation

Reported metrics are:

• Top-1 accuracy

• Top-5 accuracy

• Mean Average Precision (mAP)

The baseline method can be configured by inheriting from this class and overriding specific attributes or functions as well as passing a config object.

feature_dim = 1536

method_specific_augmentation

model

skip_embedding_training = False

class embedding_benchmark.ViTEmbedding(model: ViT_B_16Classifier)

Bases: pytorch_lightning.LightningModule

This module is used to extract features from the Vision Transformer Classifier in eval mode

forward(x: torch.Tensor) → torch.Tensor

model

class embedding_benchmark.ViT_B_16Baseline(cfg)

Bases: BaselineMethod

An abstract class that holds common code of our baseline methods.

The class runs:

• embedding training

• kNN evaluation

• linear evaluation

Reported metrics are:

• Top-1 accuracy

• Top-5 accuracy

• Mean Average Precision (mAP)

The baseline method can be configured by inheriting from this class and overriding specific attributes or functions as well as passing a config object.

get_embedding_model()

Must return a model that returns features on forward pass.

feature_dim: int = 768

method_specific_augmentation

model

class embedding_benchmark.ViT_B_16Classifier(batch_size_per_device, feature_dim, num_classes)

Bases: LinearClassifier

A fully supervised model that uses the Vision Transformer model from the torchvision library

The model uses the standard ViT_B_16 model and cross entropy (as in inherited from LinearClassifier) for training

configure_optimizers()

This optimizer is a inspired the optimizer used in the lightly benchmarks for their Vision Transformer backbones specifically the AIM Model.

model: torchvision.models.vision_transformer.VisionTransformer

embedding_benchmark.clear_cache()

embedding_benchmark.knn_predict(feature: torch.Tensor, feature_bank: torch.Tensor, feature_labels: torch.Tensor, num_classes: int, knn_k: int = 200, knn_t: float = 0.1) → torch.Tensor

[Modified version from lightly, which returns the scores instead of the predictions]

Run kNN predictions on features based on a feature bank

This method is commonly used to monitor performance of self-supervised learning methods.

The default parameters are the ones used in https://arxiv.org/pdf/1805.01978v1.pdf.

# code for kNN prediction from here: # https://colab.research.google.com/github/facebookresearch/moco/blob/colab-notebook/colab/moco_cifar10_demo.ipynb

Parameters:

• feature – Tensor with shape (B, D) for which you want predictions.

• feature_bank – Tensor of shape (D, N) of a database of features used for kNN.

• feature_labels – Labels with shape (N,) for the features in the feature_bank.

• num_classes – Number of classes (e.g. 10 for CIFAR-10).

• knn_k – Number of k neighbors used for kNN.

• knn_t – Temperature parameter to reweights similarities for kNN.

Returns:

A tensor containing the kNN scores

Examples

>>> images, targets, _ = batch
>>> feature = backbone(images).squeeze()
>>> # we recommend to normalize the features
>>> feature = F.normalize(feature, dim=1)
>>> pred_labels = knn_predict(
>>>     feature,
>>>     feature_bank,
>>>     targets_bank,
>>>     num_classes=10,
>>> )

embedding_benchmark.timing_decorator(func)

embedding_benchmark.args = None

embedding_benchmark.parser