Deep learning for NeuroImaging in Python.

Note

This page is a reference documentation. It only explains the class signature, and not how to use it. Please refer to the gallery for the big picture.

class nidl.losses.InfoNCE(temperature: float = 0.1)[source]¶

Bases: Module

Normalized temperature cross-entropy loss derived from Chen et al., ICML 2020. See https://doi.org/10.48550/arXiv.2002.05709 for details.

Parameters:: temperature : float, default=0.1

scale logits by the inverse of the temperature.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(z1: Tensor, z2: Tensor)[source]¶

Forward implementation.

Parameters:

z1 : torch.Tensor of shape (batch_size, n_features)

First embedded view.

z2 : torch.Tensor of shape (batch_size, n_features)

Second embedded view.

Returns:

loss : torch.Tensor

The InfoNCE loss computed between z1 and z2.