espnet2.lid.espnet_model.ESPnetLIDModel
espnet2.lid.espnet_model.ESPnetLIDModel
class espnet2.lid.espnet_model.ESPnetLIDModel(frontend: AbsFrontend | None, specaug: AbsSpecAug | None, normalize: AbsNormalize | None, encoder: AbsEncoder | None, pooling: AbsPooling | None, projector: AbsProjector | None, loss: AbsLoss | None, extract_feats_in_collect_stats: bool | None = None, encoder_condition: ModuleDict | None = None, pooling_condition: ModuleDict | None = None, projector_condition: ModuleDict | None = None, lang2vec_conditioning_layers: List[int] | None = None, apply_intermediate_lang2vec_loss: bool = False, apply_intermediate_lang2vec_condition: bool = True, inter_lang2vec_loss_weight: float = 0.0, cutoff_gradient_from_backbone: bool | None = True, cutoff_gradient_before_condproj: bool | None = False, shared_conditioning_proj: bool | None = False, lang2vec_dim: int = 299, lang2vec_loss: str = 'mse', lang2vec_weight: float = 0.0)
Bases: AbsESPnetModel
ESPnet LID model with optional lang2vec conditioning and auxiliary prediction.
- Parameters:
- frontend – feature extractor module
- specaug – optional SpecAugment module
- normalize – optional normalization module
- encoder – frame-level encoder (e.g., ECAPA-TDNN)
- pooling – pooling layer (e.g., AttnStatPooling)
- projector – embedding projection module (e.g., RawNet3)
- loss – main loss function (e.g., AAMSoftmax)
- extract_feats_in_collect_stats – Retained for existing LID configurations.
- encoder_condition – ModuleDict of condition module(s) for encoder layers, key is layer index, value is encoder condition module (can be shared or independent)
- pooling_condition – ModuleDict of condition module(s) for pooling layers, key is layer index, value is pooling condition module (can be shared or independent)
- projector_condition – ModuleDict of condition module(s) for projector layers, key is layer index, value is projector condition module (can be shared or independent)
- lang2vec_conditioning_layers – list of encoder layer indices to apply lang2vec conditioning
- apply_intermediate_lang2vec_loss – whether to apply auxiliary lang2vec prediction loss at intermediate layers
- apply_intermediate_lang2vec_condition – whether to apply lang2vec conditioning at intermediate layers
- inter_lang2vec_loss_weight – weight for the intermediate lang2vec loss
- cutoff_gradient_from_backbone – if True, detach the projected residual before adding it to the SSL hidden states AND freeze the conditioning projection parameters. The SSL backbone and auxiliary prediction heads remain trainable.
- cutoff_gradient_before_condproj – if True, detach the predicted lang2vec vector before projection. This alone does not freeze the projection.
- shared_conditioning_proj – if True, share conditioning projection across layers
- lang2vec_dim – Dimension of the recipe’s per-language target vectors.
- lang2vec_loss – “mse” for continuous targets (e.g. geo), or “bce” for binary targets (e.g. syntax_knn), using logits before the sigmoid.
- lang2vec_weight – Weight of the auxiliary losses, disabled by default. The classification loss receives weight 1 - lang2vec_weight on supervised batches. Empty conditioning layers and zero weight retain the ordinary LID architecture and checkpoint keys.
The classification head uses the configured LID loss. When lang2vec is enabled, a linear head predicts its targets from the final embedding, and additional heads predict the same targets at selected SSL layers. These predictions can condition the speech representation. Language metadata is used only as a training target, never as an inference input.
Reference: : Geolocation-Aware Robust Spoken Language Identification https://arxiv.org/pdf/2508.17148
Initialize internal Module state, shared by both nn.Module and ScriptModule.
collect_feats(speech: Tensor, speech_lengths: Tensor, lid_labels: Tensor | None = None, **kwargs) → Dict[str, Tensor]
compute_lang2vec_loss(prediction, target, mask=None)
Reduce auxiliary loss over rows with observed targets only.
property conditioned_layers
Return wrappers in execution order without registering duplicates.
encode_frame(feats: Tensor) → Tensor
extract_feats(speech: Tensor, speech_lengths: Tensor, return_intermediate: bool = False) → Tuple[Tensor, Tensor] | Tuple[Tensor, Tensor, Tuple[Tensor, ...]]
Extract features, optionally returning auxiliary SSL predictions.
forward(speech: Tensor, speech_lengths: Tensor, lid_labels: Tensor | None = None, extract_embd: bool = False, lang2vecs: Tensor | None = None, lang2vec_mask: Tensor | None = None, **kwargs) → Tuple[Tensor, Tensor] | Tuple[Tensor, Dict[str, Tensor], Tensor] | Tensor
Forward through encoder layers and aggregate into utterance-level feature.
- Parameters:
- speech – (Batch, samples)
- speech_lengths – (Batch,)
- extract_embd – Return embeddings and predicted language labels.
- lid_labels – (Batch, ) Integer language labels used during training.
- lang2vecs – (Batch, lang2vec_dim), auxiliary training targets.
- lang2vec_mask – (Batch, 1), optional mask for languages with auxiliary supervision. Excluded languages still contribute classification loss.
project_lang_embd(utt_level_feat: Tensor) → Tensor
