espnet2.ssl.loss.hubert.HuBERTLoss

source

class espnet2.ssl.loss.hubert.HuBERTLoss(encoder_output_size: int, num_classes: int, final_dim: int, loss_type: str = 'cross_entropy', layers: List = [-1], loss_weights: List = [1.0])

Bases: AbsSSLLoss

HuBERT MLM Loss

• Parameters:
  • encoder_output_size (int) – input dimension
  • num_classes (int) – vocab size
  • final_dim (int) – final projection dim
  • loss_type (str) – TODO, unused for now
  • layers (List) – encoder output layers for loss
  • loss_weights (List) – weight of each layer for loss

forward(encoder_output: List, encoder_output_lengths: Tensor = None, text: Tensor = None, text_lengths: Tensor = None, mask_info: Dict = None) → Tuple[Tensor, Dict]

HuBERT forward

• Parameters:
  • encoder_output (List) – List of encoded sequences (B, T, D) from each layer.
  • encoder_output_lengths (Tensor) – Lengths of batched encoder sequences (B,).
  • text (Tensor) – text targets (B, T)
  • text_lengths (Tensor) – Lengths of text targets (B,).
  • mask_info (Dict) – Contains masked/unmasked indices