CLIPLoss#

class core.nn.loss.CLIPLoss(temperature=0.5, gather_distributed=True)[source]#

Bases: torch.nn.modules.module.Module

Symmetric InfoNCE loss over two paired views (CLIP, Eq. 1 of the NEMO paper).

Parameters:

  • temperature (float) – softmax temperature applied to the cosine similarities.

  • gather_distributed (bool) – take the negatives from the global batch by pooling both views across ranks, so the objective does not depend on the GPU count. Every rank must hold the same number of samples.

forward(z_a, z_b)[source]#

Compute the loss for row-wise paired embeddings.

Parameters:

  • z_a (Tensor) – \((N, D)\) float, embeddings of one view.

  • z_b (Tensor) – \((N, D)\) float, embeddings of the other view.

Return type:

Tensor