# Copyright 2021 Tianzi Wang
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0
# Thanks to Abdelrahman Mohamed and Wei-Ning Hsu's help in this implementation,
# Their origial Hubert work is in:
# Paper: https://arxiv.org/pdf/2106.07447.pdf
# Code in Fairseq: https://github.com/pytorch/fairseq/tree/master/examples/hubert
"""Encoder definition."""
import contextlib
import copy
import logging
import os
from pathlib import Path
from typing import List, Optional, Tuple
import torch
import yaml
from filelock import FileLock
from typeguard import typechecked
from espnet2.asr.encoder.abs_encoder import AbsEncoder
from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
from espnet.nets.pytorch_backend.transformer.layer_norm import LayerNorm
[docs]class TorchAudioHuBERTPretrainEncoder(AbsEncoder):
"""Torch Audio Hubert encoder module.
Args:
extractor_mode: Operation mode of feature extractor.
Valid values are "group_norm" or "layer_norm".
extractor_conv_layer_config: Configuration of convolution layers in feature
extractor. List of convolution configuration,
i.e. [[output_channel, kernel_size, stride], ...]
extractor_conv_bias: Whether to include bias term to each convolution
operation.
encoder_embed_dim: The dimension of embedding in encoder.
encoder_projection_dropout: The dropout probability applied after the input
feature is projected to "encoder_embed_dim".
encoder_pos_conv_kernel: Kernel size of convolutional positional embeddings.
encoder_pos_conv_groups: Number of groups of convolutional positional
embeddings.
encoder_num_layers: Number of self attention layers in transformer block.
encoder_num_heads: Number of heads in self attention layers.
encoder_attention_dropout: Dropout probability applied after softmax in
self-attention layer.
encoder_ff_interm_features: Dimension of hidden features in feed forward layer.
encoder_ff_interm_dropout: Dropout probability applied in feedforward layer.
encoder_dropout: Dropout probability applied at the end of feed forward layer.
encoder_layer_norm_first: Control the order of layer norm in transformer layer
and each encoder layer. If True, in transformer layer, layer norm is
applied before features are fed to encoder layers.
encoder_layer_drop: Probability to drop each encoder layer during training.
mask_prob: Probability for each token to be chosen as start of the span
to be masked.
mask_selection: How to choose the mask length.
Options: [static, uniform, normal, poisson].
mask_other: Secondary mask argument (used for more complex distributions).
mask_length: The lengths of the mask.
no_mask_overlap: Whether to allow masks to overlap.
mask_min_space: Minimum space between spans (if no overlap is enabled).
mask_channel_prob: (float): The probability of replacing a feature with 0.
mask_channel_selection: How to choose the mask length for channel masking.
Options: [static, uniform, normal, poisson].
mask_channel_other: Secondary mask argument for channel masking(used for more
complex distributions).
mask_channel_length: Minimum space between spans (if no overlap is enabled)
for channel masking.
no_mask_channel_overlap: Whether to allow channel masks to overlap.
mask_channel_min_space: Minimum space between spans for channel
masking(if no overlap is enabled).
skip_masked: If True, skip computing losses over masked frames.
skip_nomask: If True, skip computing losses over unmasked frames.
num_classes: The number of classes in the labels.
final_dim: Project final representations and targets to final_dim.
feature_grad_mult: The factor to scale the convolutional feature extraction
layer gradients by. The scale factor will not affect the forward pass.
finetuning: Whether to finetuning the model with ASR or other tasks.
freeze_encoder_updates: The number of steps to freeze the encoder parameters
in ASR finetuning.
Hubert specific Args:
Please refer to:
https://pytorch.org/audio/stable/generated/torchaudio.models.hubert_pretrain_model.html#torchaudio.models.hubert_pretrain_model
"""
@typechecked
def __init__(
self,
input_size: int = None,
extractor_mode: str = "group_norm",
extractor_conv_layer_config: Optional[List[List[int]]] = [
[512, 10, 5],
[512, 3, 2],
[512, 3, 2],
[512, 3, 2],
[512, 3, 2],
[512, 2, 2],
[512, 2, 2],
],
extractor_conv_bias: bool = False,
encoder_embed_dim: int = 768,
encoder_projection_dropout: float = 0.1,
encoder_pos_conv_kernel: int = 128,
encoder_pos_conv_groups: int = 16,
encoder_num_layers: int = 12,
encoder_num_heads: int = 12,
encoder_attention_dropout: float = 0.1,
encoder_ff_interm_features: int = 3072,
encoder_ff_interm_dropout: float = 0.0,
encoder_dropout: float = 0.1,
encoder_layer_norm_first: bool = False,
encoder_layer_drop: float = 0.05,
mask_prob: float = 0.8,
mask_selection: str = "static",
mask_other: float = 0.0,
mask_length: int = 10,
no_mask_overlap: bool = False,
mask_min_space: int = 1,
mask_channel_prob: float = 0.0,
mask_channel_selection: str = "static",
mask_channel_other: float = 0.0,
mask_channel_length: int = 10,
no_mask_channel_overlap: bool = False,
mask_channel_min_space: int = 1,
skip_masked: bool = False,
skip_nomask: bool = False,
num_classes: int = 100,
final_dim: int = 256,
feature_grad_mult: Optional[float] = 0.1,
finetuning: bool = False,
freeze_encoder_updates: int = 0,
):
super().__init__()
try:
import torchaudio
except Exception as e:
print("Error: torchaudio is not properly installed.")
print("Please install torchaudio")
raise e
self._output_size = encoder_embed_dim
self.hubert_pretrain_model = torchaudio.models.hubert_pretrain_model(
extractor_mode=extractor_mode,
extractor_conv_layer_config=extractor_conv_layer_config,
extractor_conv_bias=extractor_conv_bias,
encoder_embed_dim=encoder_embed_dim,
encoder_projection_dropout=encoder_projection_dropout,
encoder_pos_conv_kernel=encoder_pos_conv_kernel,
encoder_pos_conv_groups=encoder_pos_conv_groups,
encoder_num_layers=encoder_num_layers,
encoder_num_heads=encoder_num_heads,
encoder_attention_dropout=encoder_attention_dropout,
encoder_ff_interm_features=encoder_ff_interm_features,
encoder_ff_interm_dropout=encoder_ff_interm_dropout,
encoder_dropout=encoder_dropout,
encoder_layer_norm_first=encoder_layer_norm_first,
encoder_layer_drop=encoder_layer_drop,
mask_prob=mask_prob,
mask_selection=mask_selection,
mask_other=mask_other,
mask_length=mask_length,
no_mask_overlap=no_mask_overlap,
mask_min_space=mask_min_space,
mask_channel_prob=mask_channel_prob,
mask_channel_selection=mask_channel_selection,
mask_channel_other=mask_channel_other,
mask_channel_length=mask_channel_length,
no_mask_channel_overlap=no_mask_channel_overlap,
mask_channel_min_space=mask_channel_min_space,
skip_masked=skip_masked,
skip_nomask=skip_nomask,
num_classes=num_classes,
final_dim=final_dim,
feature_grad_mult=feature_grad_mult,
)
self.pretrained_params = copy.deepcopy(self.hubert_pretrain_model.state_dict())
self.finetuning = finetuning
if finetuning:
for p in self.hubert_pretrain_model.wav2vec2.feature_extractor.parameters():
p.requires_grad = False
self.register_buffer("global_step", torch.LongTensor([0]))
self.freeze_encoder_updates = freeze_encoder_updates
[docs] def output_size(self) -> int:
return self._output_size
[docs] def forward(
self,
xs_pad: torch.Tensor,
ilens: torch.Tensor,
ys_pad: torch.Tensor = None,
ys_pad_length: torch.Tensor = None,
prev_states: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
"""Forward Hubert Pretrain Encoder.
Args:
xs_pad: input tensor (B, L, D)
ilens: input length (B)
prev_states: Not to be used now.
Returns:
position embedded tensor and mask
"""
if not self.finetuning:
return self._pretraining_forward(xs_pad, ilens, ys_pad)
else:
if self.training:
return self._finetuning_forward(xs_pad, ilens)
else:
return self._eval_forward(xs_pad, ilens)
def _pretraining_forward(self, xs_pad, ilens, ys_pad):
assert ys_pad is not None
(
logit_m,
logit_u,
feature_penalty,
) = self.hubert_pretrain_model.forward(xs_pad, ys_pad, ilens)
return logit_m, logit_u, feature_penalty
def _finetuning_forward(self, xs_pad, ilens):
def get_padding_mask(input, lengths):
"""get_padding_mask() from torchaudio.models.wav2vec2.components"""
batch_size, max_len, _ = input.shape
mask = (
torch.arange(max_len, device=lengths.device).expand(batch_size, max_len)
>= lengths[:, None]
)
return mask
# manually add the steps. It is not accurate.
# TODO(simpleoier): to introduce the global update steps into encoder module
self.global_step += 1
if self.global_step <= self.freeze_encoder_updates:
with torch.no_grad():
x, out_len = self.hubert_pretrain_model.wav2vec2.feature_extractor(
xs_pad, ilens
)
padding_mask = get_padding_mask(x, out_len)
(
x,
attention_mask,
) = self.hubert_pretrain_model.wav2vec2.encoder._preprocess(x, out_len)
x, _ = self.hubert_pretrain_model.mask_generator(x, padding_mask)
x = self.hubert_pretrain_model.wav2vec2.encoder.transformer(
x, attention_mask=attention_mask
)
else:
with torch.no_grad():
x, out_len = self.hubert_pretrain_model.wav2vec2.feature_extractor(
xs_pad, ilens
)
padding_mask = get_padding_mask(x, out_len)
(
x,
attention_mask,
) = self.hubert_pretrain_model.wav2vec2.encoder._preprocess(x, out_len)
x, _ = self.hubert_pretrain_model.mask_generator(x, padding_mask)
x = self.hubert_pretrain_model.wav2vec2.encoder.transformer(
x, attention_mask=attention_mask
)
return x, (~padding_mask).long().sum(dim=1), None
def _eval_forward(self, xs_pad, ilens):
x, lengths = self.hubert_pretrain_model.wav2vec2.feature_extractor(
xs_pad, ilens
)
x = self.hubert_pretrain_model.wav2vec2.encoder(x, lengths)
return x, lengths, None
[docs] def reload_pretrained_parameters(self):
self.hubert_pretrain_model.load_state_dict(self.pretrained_params, strict=False)
logging.info("Pretrained Hubert model parameters reloaded!")
[docs]class FairseqHubertEncoder(AbsEncoder):
"""FairSeq Hubert encoder module, used for loading pretrained weight and finetuning
Args:
input_size: input dim
hubert_url: url to Hubert pretrained model
hubert_dir_path: directory to download the Wav2Vec2.0 pretrained model.
output_size: dimension of attention
normalize_before: whether to use layer_norm before the first block
freeze_finetune_updates: steps that freeze all layers except output layer
before tuning the whole model (nessasary to prevent overfit).
dropout_rate: dropout rate
activation_dropout: dropout rate in activation function
attention_dropout: dropout rate in attention
Hubert specific Args:
Please refer to:
https://github.com/pytorch/fairseq/blob/master/fairseq/models/hubert/hubert.py
"""
@typechecked
def __init__(
self,
input_size: int,
hubert_url: str = "./",
hubert_dir_path: str = "./",
output_size: int = 256,
normalize_before: bool = False,
freeze_finetune_updates: int = 0,
dropout_rate: float = 0.0,
activation_dropout: float = 0.1,
attention_dropout: float = 0.0,
mask_length: int = 10,
mask_prob: float = 0.75,
mask_selection: str = "static",
mask_other: int = 0,
apply_mask: bool = True,
mask_channel_length: int = 64,
mask_channel_prob: float = 0.5,
mask_channel_other: int = 0,
mask_channel_selection: str = "static",
layerdrop: float = 0.1,
feature_grad_mult: float = 0.0,
):
super().__init__()
self.apply_mask = apply_mask
try:
import fairseq
from fairseq.models.hubert.hubert import HubertModel
except Exception as e:
print("Error: FairSeq is not properly installed.")
print("Please install FairSeq: cd ${MAIN_ROOT}/tools && make fairseq.done")
raise e
arg_overrides = {
"dropout": dropout_rate,
"activation_dropout": activation_dropout,
"attention_dropout": attention_dropout,
"mask_length": mask_length,
"mask_prob": mask_prob,
"mask_selection": mask_selection,
"mask_other": mask_other,
"mask_channel_length": mask_channel_length,
"mask_channel_prob": mask_channel_prob,
"mask_channel_selection": mask_channel_selection,
"mask_channel_other": mask_channel_other,
"encoder_layerdrop": layerdrop,
"feature_grad_mult": feature_grad_mult,
"data": hubert_dir_path,
}
if hubert_url == "espnet":
self.hubert_model_path = hubert_dir_path
s = torch.load(
self.hubert_model_path,
map_location=torch.device("cpu"),
)
if all("encoder.encoder" in k for k in s):
try:
state = {
k.replace("encoder.encoder.", ""): v
for k, v in s.items()
if "label_embs_concat" not in k
}
except Exception as e:
raise e
config_file = os.path.join(
"/".join(self.hubert_model_path.split("/")[:-1]),
"config.yaml",
)
config_file = Path(config_file)
with config_file.open("r", encoding="utf-8") as f:
self.pretrained_cfg = yaml.safe_load(f)
model = FairseqHubertPretrainEncoder(
input_size=self.pretrained_cfg["input_size"],
hubert_dict=self.pretrained_cfg["hubert_dict"],
**self.pretrained_cfg["encoder_conf"],
)
model = model.encoder
d = self.pretrained_cfg["encoder_conf"]["output_size"]
self.pretrained_params = copy.deepcopy(state)
else:
self.hubert_model_path = download_hubert(hubert_url, hubert_dir_path)
(
models,
self.pretrained_cfg,
task,
) = fairseq.checkpoint_utils.load_model_ensemble_and_task(
[self.hubert_model_path],
arg_overrides=arg_overrides,
strict=False,
)
model = models[0]
d = self.pretrained_cfg.model.encoder_embed_dim
self.pretrained_params = copy.deepcopy(model.state_dict())
self._output_size = output_size
if not isinstance(model, HubertModel):
try:
model = model.hubert_encoder.hubert_model
except Exception as e:
print(
"Error: pretrained models should be within: "
"'HubertModel, Hubertctc' classes, etc."
)
raise e
self.encoders = model
self.normalize_before = normalize_before
if self.normalize_before:
self.after_norm = LayerNorm(output_size)
if output_size and output_size != d:
self.output_layer = torch.nn.Sequential(
torch.nn.Linear(d, output_size),
)
else:
self.output_layer = None
self.freeze_finetune_updates = freeze_finetune_updates
self.register_buffer("num_updates", torch.LongTensor([0]))
[docs] def output_size(self) -> int:
return self._output_size
[docs] def forward(
self,
xs_pad: torch.Tensor,
ilens: torch.Tensor,
prev_states: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
"""Forward Hubert ASR Encoder.
Args:
xs_pad: input tensor (B, L, D)
ilens: input length (B)
prev_states: Not to be used now.
Returns:
position embedded tensor and mask
"""
masks = make_pad_mask(ilens).to(xs_pad.device)
ft = self.freeze_finetune_updates <= self.num_updates
if self.num_updates <= self.freeze_finetune_updates:
self.num_updates += 1
elif ft and self.num_updates == self.freeze_finetune_updates + 1:
self.num_updates += 1
logging.info("Start fine-tuning hubert parameters!")
else:
self.num_updates += 1
with torch.no_grad() if not ft else contextlib.nullcontext():
enc_outputs = self.encoders(
xs_pad,
padding_mask=masks,
mask=self.apply_mask and self.training,
features_only=True,
output_layer=None,
)
xs_pad = enc_outputs["x"] # (B,T,C),
masks = enc_outputs["padding_mask"] # (B, T)
# save gpu memory
del enc_outputs
olens = (~masks).sum(dim=1)
if self.output_layer is not None:
xs_pad = self.output_layer(xs_pad)
if self.normalize_before:
xs_pad = self.after_norm(xs_pad)
return xs_pad, olens, None
[docs] def reload_pretrained_parameters(self):
self.encoders.load_state_dict(self.pretrained_params, strict=False)
logging.info("Pretrained Hubert model parameters reloaded!")
[docs]class FairseqHubertPretrainEncoder(AbsEncoder):
"""FairSeq Hubert pretrain encoder module, only used for pretraining stage
Args:
input_size: input dim
output_size: dimension of attention
linear_units: dimension of feedforward layers
attention_heads: the number of heads of multi head attention
num_blocks: the number of encoder blocks
dropout_rate: dropout rate
attention_dropout_rate: dropout rate in attention
hubert_dict: target dictionary for Hubert pretraining
label_rate: label frame rate. -1 for sequence label
sample_rate: target sample rate.
use_amp: whether to use automatic mixed precision
normalize_before: whether to use layer_norm before the first block
"""
@typechecked
def __init__(
self,
input_size: int = 1,
output_size: int = 1024,
linear_units: int = 1024,
attention_heads: int = 12,
num_blocks: int = 12,
dropout_rate: float = 0.0,
attention_dropout_rate: float = 0.0,
activation_dropout_rate: float = 0.0,
hubert_dict: str = "./dict.txt",
label_rate: int = 100,
checkpoint_activations: bool = False,
sample_rate: int = 16000,
use_amp: bool = False,
**kwargs,
):
super().__init__()
self._output_size = output_size
self.use_amp = use_amp
try:
from fairseq.data.dictionary import Dictionary
from fairseq.models.hubert.hubert import HubertConfig # noqa: H301
from fairseq.models.hubert.hubert import HubertModel # noqa: H301
from fairseq.models.hubert.hubert import ( # noqa: H301
HubertPretrainingConfig,
)
except Exception as e:
print("Error: FairSeq is not properly installed.")
print("Please install FairSeq: cd ${MAIN_ROOT}/tools && make fairseq.done")
raise e
cfg_overides = {
"encoder_embed_dim": output_size,
"encoder_ffn_embed_dim": linear_units,
"encoder_attention_heads": attention_heads,
"encoder_layers": num_blocks,
"final_dim": output_size,
"dropout": dropout_rate,
"attention_dropout": attention_dropout_rate,
"label_rate": label_rate,
"checkpoint_activations": checkpoint_activations,
}
cfg_overides = {**cfg_overides, **kwargs}
self.cfg = HubertConfig()
for key, value in cfg_overides.items():
if hasattr(self.cfg, key):
setattr(self.cfg, key, value)
hubert_task_cfg = HubertPretrainingConfig()
hubert_task_cfg_overides = {
"label_rate": label_rate,
"sample_rate": sample_rate,
}
for key, value in hubert_task_cfg_overides.items():
if hasattr(hubert_task_cfg, key):
setattr(hubert_task_cfg, key, value)
d = Dictionary()
self._build_dictionary(d, hubert_dict)
self.encoder = HubertModel(self.cfg, hubert_task_cfg, self.dictionaries)
def _build_dictionary(self, dictionary, hubert_dict_path):
if os.path.exists(f"{hubert_dict_path}"):
setattr(dictionary, "symbols", [])
setattr(dictionary, "count", [])
setattr(dictionary, "indices", {})
dictionary.add_from_file(f"{hubert_dict_path}")
else:
dictionary.add_symbol("0")
self.dictionaries = [dictionary]
[docs] def output_size(self) -> int:
return self._output_size
[docs] def forward(
self,
xs_pad: torch.Tensor,
ilens: torch.Tensor,
ys_pad: torch.Tensor,
ys_pad_length: torch.Tensor,
prev_states: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
"""Forward Hubert Pretrain Encoder.
Args:
xs_pad: input tensor (B, L, D)
ilens: input length (B)
prev_states: Not to be used now.
Returns:
position embedded tensor and mask
"""
self.cast_mask_emb()
masks = make_pad_mask(ilens).to(xs_pad.device)
ys_pad = ys_pad[:, : min(ys_pad_length)]
enc_outputs = self.encoder(
xs_pad,
padding_mask=masks,
mask=True,
target_list=[ys_pad],
features_only=False,
)
return enc_outputs
[docs] def cast_mask_emb(self):
if self.use_amp and self.encoder.mask_emb.dtype != torch.cuda.HalfTensor:
self.encoder.mask_emb = torch.nn.Parameter(self.encoder.mask_emb.half())
[docs] def reload_pretrained_parameters(self):
self.encoder.mask_emb = torch.nn.Parameter(
torch.HalfTensor(self.cfg.encoder_embed_dim).uniform_()
)
logging.info(
f"Hubert mask embedding re-initiallized!, \
{self.encoder.mask_emb.dtype}, \
{self.use_amp}"
)
[docs]def download_hubert(model_url, dir_path):
os.makedirs(dir_path, exist_ok=True)
model_name = model_url.split("/")[-1]
model_path = os.path.join(dir_path, model_name)
with FileLock(model_path + ".lock"):
if not os.path.exists(model_path):
torch.hub.download_url_to_file(model_url, model_path)
logging.info(f"Hubert model downloaded {model_path}")
else:
logging.info(f"Hubert model {model_path} already exists.")
return model_path