quaterion.loss.arcface_loss module

class ArcFaceLoss(embedding_size: int, num_groups: int, scale: float = 64.0, margin: float = 0.5)[source]

Bases: GroupLoss

Additive Angular Margin Loss as defined in https://arxiv.org/abs/1801.07698

Parameters:

• embedding_size – Output dimension of the encoder.

• num_groups – Number of groups in the dataset.

• scale – Scaling value to make cross entropy work.

• margin – Margin value to push groups apart.

forward(embeddings, groups) → Tensor[source]

Compute loss value

Parameters:

• embeddings – shape: (batch_size, vector_length) - Output embeddings from the encoder.

• groups – shape: (batch_size,) - Group ids associated with embeddings.

Returns:

Tensor – loss value.

training: bool

l2_norm(inputs: Tensor, dim: int = 0) → Tensor[source]

Apply L2 normalization to tensor

Parameters:

• inputs – Input tensor.

• dim – Dimension to operate on.

Returns:

torch.Tensor – L2-normalized tensor