ContrastiveHead

class mmpretrain.models.heads.ContrastiveHead(loss, temperature=0.1, init_cfg=None)

Head for contrastive learning.

The contrastive loss is implemented in this head and is used in SimCLR, MoCo, DenseCL, etc.

Parameters:

• loss (dict) – Config dict for module of loss functions.

• temperature (float) – The temperature hyper-parameter that controls the concentration level of the distribution. Defaults to 0.1.

• init_cfg (dict or List[dict], optional) – Initialization config dict. Defaults to None.

loss(pos, neg)

Forward function to compute contrastive loss.

Parameters:

• pos (torch.Tensor) – Nx1 positive similarity.

• neg (torch.Tensor) – Nxk negative similarity.

Returns:

The contrastive loss.

Return type:

torch.Tensor