JointEmbeddingTrainer

class stable_ssl.trainers.JointEmbeddingTrainer(data, module, hardware, optim, logger, loss=None, **kwargs)

Bases: BaseTrainer

Base class for training a joint-embedding SSL model.

compute_loss()

Compute final loss as sum of SSL loss and classifier losses.

compute_loss_classifiers(representations, embeddings, labels)

Compute the classifier loss for both backbone and projector.

forward(*args, **kwargs)

Forward pass. By default, it simply calls the ‘backbone’ module.

predict()

Call the backbone classifier on the forward pass of current batch.