MoCov3

MoCov3#

class stable_pretraining.methods.MoCov3(encoder_name: str | Module = 'vit_small_patch16_224', projector_dims: Sequence[int] = (4096, 4096, 256), predictor_hidden_dim: int = 4096, temperature: float = 0.2, ema_decay_start: float = 0.99, ema_decay_end: float = 1.0, low_resolution: bool = False, pretrained: bool = False)[source]#

Bases: Module

MoCo v3: ViT-friendly momentum contrastive learning.

Architecture:

  • Backbone (student) wrapped with EMA teacher.

  • Projector (student) wrapped with EMA teacher.

  • Predictor on the student side only.

Parameters:

  • encoder_name – timm model or pre-built nn.Module.

  • projector_dims – 3-layer projector dims (default (4096, 4096, 256)).

  • predictor_hidden_dim – Predictor hidden dim (default 4096).

  • temperature – InfoNCE temperature (default 0.2).

  • ema_decay_start – Initial EMA (default 0.99).

  • ema_decay_end – Final EMA (default 1.0).

  • low_resolution – Adapt first conv for low-res input.

  • pretrained – Load pretrained timm weights.

forward(view1: Tensor, view2: Tensor | None = None) → MoCov3Output[source]#

Same as torch.nn.Module.forward().

Parameters:

  • *args – Whatever you decide to pass into the forward method.

  • **kwargs – Keyword arguments are also possible.

Returns:

Your model’s output