wehub/modelscope--funasr
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:25:10 +08:00
commit c397331b1e
3684 changed files with 990993 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
funasr/models/data2vec/__init__.py
View File
funasr/models/data2vec/data2vec.py
+185
View File
@@ -0,0 +1,185 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from contextlib import contextmanager
from distutils.version import LooseVersion
from typing import Dict
from typing import Optional
from typing import Tuple
import torch
import torch.nn as nn
# from funasr.layers.abs_normalize import AbsNormalize
# from funasr.models.base_model import FunASRModel
# from funasr.models.encoder.abs_encoder import AbsEncoder
from funasr.frontends.abs_frontend import AbsFrontend
# from funasr.models.preencoder.abs_preencoder import AbsPreEncoder
# from funasr.models.specaug.abs_specaug import AbsSpecAug
from funasr.train_utils.device_funcs import force_gatherable
if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
from torch.cuda.amp import autocast
else:
# Nothing to do if torch<1.6.0
@contextmanager
def autocast(enabled=True):
"""Autocast.
Args:
enabled: TODO.
"""
yield
class Data2VecPretrainModel(nn.Module):
"""Data2Vec Pretrain model"""
def __init__(
self,
frontend=None,
specaug=None,
normalize=None,
encoder=None,
preencoder=None,
):
"""Initialize Data2VecPretrainModel.
Args:
frontend: Audio frontend for feature extraction.
specaug: TODO.
normalize: TODO.
encoder: TODO.
preencoder: TODO.
"""
super().__init__()
self.frontend = frontend
self.specaug = specaug
self.normalize = normalize
self.preencoder = preencoder
self.encoder = encoder
self.num_updates = 0
def forward(
self,
speech: torch.Tensor,
speech_lengths: torch.Tensor,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
"""Frontend + Encoder + Calc loss
Args:
speech: (Batch, Length, ...)
speech_lengths: (Batch, )
"""
# Check that batch_size is unified
assert speech.shape[0] == speech_lengths.shape[0], (speech.shape, speech_lengths.shape)
self.encoder.set_num_updates(self.num_updates)
# 1. Encoder
encoder_out = self.encode(speech, speech_lengths)
losses = encoder_out["losses"]
loss = sum(losses.values())
sample_size = encoder_out["sample_size"]
loss = loss.sum() / sample_size
target_var = float(encoder_out["target_var"])
pred_var = float(encoder_out["pred_var"])
ema_decay = float(encoder_out["ema_decay"])
stats = dict(
loss=torch.clone(loss.detach()),
target_var=target_var,
pred_var=pred_var,
ema_decay=ema_decay,
)
loss, stats, weight = force_gatherable((loss, stats, sample_size), loss.device)
return loss, stats, weight
def collect_feats(
self, speech: torch.Tensor, speech_lengths: torch.Tensor
) -> Dict[str, torch.Tensor]:
"""Collect feats.
Args:
speech: Speech audio tensor, shape (batch, time).
speech_lengths: Length of each speech sample.
"""
feats, feats_lengths = self._extract_feats(speech, speech_lengths)
return {"feats": feats, "feats_lengths": feats_lengths}
def encode(
self,
speech: torch.Tensor,
speech_lengths: torch.Tensor,
):
"""Frontend + Encoder.
Args:
speech: (Batch, Length, ...)
speech_lengths: (Batch, )
"""
with autocast(False):
# 1. Extract feats
feats, feats_lengths = self._extract_feats(speech, speech_lengths)
# 2. Data augmentation
if self.specaug is not None and self.training:
feats, feats_lengths = self.specaug(feats, feats_lengths)
# 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
if self.normalize is not None:
feats, feats_lengths = self.normalize(feats, feats_lengths)
# Pre-encoder, e.g. used for raw input data
if self.preencoder is not None:
feats, feats_lengths = self.preencoder(feats, feats_lengths)
# 4. Forward encoder
if min(speech_lengths) == max(speech_lengths): # for clipping, set speech_lengths as None
speech_lengths = None
encoder_out = self.encoder(feats, speech_lengths, mask=True, features_only=False)
return encoder_out
def _extract_feats(
self, speech: torch.Tensor, speech_lengths: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Internal: extract feats.
Args:
speech: Speech audio tensor, shape (batch, time).
speech_lengths: Length of each speech sample.
"""
assert speech_lengths.dim() == 1, speech_lengths.shape
# for data-parallel
speech = speech[:, : speech_lengths.max()]
if self.frontend is not None:
# Frontend
# e.g. STFT and Feature extract
# data_loader may send time-domain signal in this case
# speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim)
feats, feats_lengths = self.frontend(speech, speech_lengths)
else:
# No frontend and no feature extract
feats, feats_lengths = speech, speech_lengths
return feats, feats_lengths
def set_num_updates(self, num_updates):
"""Set num updates.
Args:
num_updates: TODO.
"""
self.num_updates = num_updates
def get_num_updates(self):
"""Get num updates."""
return self.num_updates
funasr/models/data2vec/data2vec_encoder.py
+680
View File
@@ -0,0 +1,680 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import logging
import math
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from funasr.models.data2vec.data_utils import compute_mask_indices
from funasr.models.data2vec.ema_module import EMAModule
from funasr.models.data2vec.grad_multiply import GradMultiply
from funasr.models.data2vec.wav2vec2 import (
ConvFeatureExtractionModel,
TransformerEncoder,
)
from funasr.models.transformer.utils.nets_utils import make_pad_mask
def get_annealed_rate(start, end, curr_step, total_steps):
"""Get annealed rate.
Args:
start: TODO.
end: TODO.
curr_step: TODO.
total_steps: TODO.
"""
r = end - start
pct_remaining = 1 - curr_step / total_steps
return end - r * pct_remaining
class Data2VecEncoder(nn.Module):
def __init__(
self,
# for ConvFeatureExtractionModel
input_size: int = None,
extractor_mode: str = None,
conv_feature_layers: str = "[(512,2,2)] + [(512,2,2)]",
# for Transformer Encoder
## model architecture
layer_type: str = "transformer",
layer_norm_first: bool = False,
encoder_layers: int = 12,
encoder_embed_dim: int = 768,
encoder_ffn_embed_dim: int = 3072,
encoder_attention_heads: int = 12,
activation_fn: str = "gelu",
## dropouts
dropout: float = 0.1,
attention_dropout: float = 0.1,
activation_dropout: float = 0.0,
encoder_layerdrop: float = 0.0,
dropout_input: float = 0.0,
dropout_features: float = 0.0,
## grad settings
feature_grad_mult: float = 1.0,
## masking
mask_prob: float = 0.65,
mask_length: int = 10,
mask_selection: str = "static",
mask_other: int = 0,
no_mask_overlap: bool = False,
mask_min_space: int = 1,
require_same_masks: bool = True, # if set as True, collate_fn should be clipping
mask_dropout: float = 0.0,
## channel masking
mask_channel_length: int = 10,
mask_channel_prob: float = 0.0,
mask_channel_before: bool = False,
mask_channel_selection: str = "static",
mask_channel_other: int = 0,
no_mask_channel_overlap: bool = False,
mask_channel_min_space: int = 1,
## positional embeddings
conv_pos: int = 128,
conv_pos_groups: int = 16,
pos_conv_depth: int = 1,
max_positions: int = 100000,
# EMA module
average_top_k_layers: int = 8,
layer_norm_target_layer: bool = False,
instance_norm_target_layer: bool = False,
instance_norm_targets: bool = False,
layer_norm_targets: bool = False,
batch_norm_target_layer: bool = False,
group_norm_target_layer: bool = False,
ema_decay: float = 0.999,
ema_end_decay: float = 0.9999,
ema_anneal_end_step: int = 100000,
ema_transformer_only: bool = True,
ema_layers_only: bool = True,
min_target_var: float = 0.1,
min_pred_var: float = 0.01,
# Loss
loss_beta: float = 0.0,
loss_scale: float = None,
# FP16 optimization
required_seq_len_multiple: int = 2,
):
"""Initialize Data2VecEncoder.
Args:
input_size: Size/dimension parameter.
extractor_mode: TODO.
conv_feature_layers: TODO.
layer_type: TODO.
layer_norm_first: TODO.
encoder_layers: TODO.
encoder_embed_dim: Size/dimension parameter.
encoder_ffn_embed_dim: Size/dimension parameter.
encoder_attention_heads: TODO.
activation_fn: TODO.
dropout: TODO.
attention_dropout: TODO.
activation_dropout: TODO.
encoder_layerdrop: TODO.
dropout_input: TODO.
dropout_features: TODO.
feature_grad_mult: TODO.
mask_prob: TODO.
mask_length: TODO.
mask_selection: TODO.
mask_other: TODO.
no_mask_overlap: TODO.
mask_min_space: TODO.
require_same_masks: TODO.
mask_dropout: TODO.
mask_channel_length: TODO.
mask_channel_prob: TODO.
mask_channel_before: TODO.
mask_channel_selection: TODO.
mask_channel_other: TODO.
no_mask_channel_overlap: TODO.
mask_channel_min_space: TODO.
conv_pos: TODO.
conv_pos_groups: TODO.
pos_conv_depth: TODO.
max_positions: TODO.
average_top_k_layers: TODO.
layer_norm_target_layer: TODO.
instance_norm_target_layer: TODO.
instance_norm_targets: TODO.
layer_norm_targets: TODO.
batch_norm_target_layer: TODO.
group_norm_target_layer: TODO.
ema_decay: TODO.
ema_end_decay: TODO.
ema_anneal_end_step: TODO.
ema_transformer_only: TODO.
ema_layers_only: TODO.
min_target_var: TODO.
min_pred_var: TODO.
loss_beta: TODO.
loss_scale: TODO.
required_seq_len_multiple: TODO.
"""
super().__init__()
# ConvFeatureExtractionModel
self.conv_feature_layers = conv_feature_layers
feature_enc_layers = eval(conv_feature_layers)
self.extractor_embed = feature_enc_layers[-1][0]
self.feature_extractor = ConvFeatureExtractionModel(
conv_layers=feature_enc_layers,
dropout=0.0,
mode=extractor_mode,
in_d=input_size,
)
# Transformer Encoder
## model architecture
self.layer_type = layer_type
self.layer_norm_first = layer_norm_first
self.encoder_layers = encoder_layers
self.encoder_embed_dim = encoder_embed_dim
self.encoder_ffn_embed_dim = encoder_ffn_embed_dim
self.encoder_attention_heads = encoder_attention_heads
self.activation_fn = activation_fn
## dropout
self.dropout = dropout
self.attention_dropout = attention_dropout
self.activation_dropout = activation_dropout
self.encoder_layerdrop = encoder_layerdrop
self.dropout_input = dropout_input
self.dropout_features = dropout_features
## grad settings
self.feature_grad_mult = feature_grad_mult
## masking
self.mask_prob = mask_prob
self.mask_length = mask_length
self.mask_selection = mask_selection
self.mask_other = mask_other
self.no_mask_overlap = no_mask_overlap
self.mask_min_space = mask_min_space
self.require_same_masks = (
require_same_masks # if set as True, collate_fn should be clipping
)
self.mask_dropout = mask_dropout
## channel masking
self.mask_channel_length = mask_channel_length
self.mask_channel_prob = mask_channel_prob
self.mask_channel_before = mask_channel_before
self.mask_channel_selection = mask_channel_selection
self.mask_channel_other = mask_channel_other
self.no_mask_channel_overlap = no_mask_channel_overlap
self.mask_channel_min_space = mask_channel_min_space
## positional embeddings
self.conv_pos = conv_pos
self.conv_pos_groups = conv_pos_groups
self.pos_conv_depth = pos_conv_depth
self.max_positions = max_positions
self.mask_emb = nn.Parameter(torch.FloatTensor(self.encoder_embed_dim).uniform_())
self.encoder = TransformerEncoder(
dropout=self.dropout,
encoder_embed_dim=self.encoder_embed_dim,
required_seq_len_multiple=required_seq_len_multiple,
pos_conv_depth=self.pos_conv_depth,
conv_pos=self.conv_pos,
conv_pos_groups=self.conv_pos_groups,
# transformer layers
layer_type=self.layer_type,
encoder_layers=self.encoder_layers,
encoder_ffn_embed_dim=self.encoder_ffn_embed_dim,
encoder_attention_heads=self.encoder_attention_heads,
attention_dropout=self.attention_dropout,
activation_dropout=self.activation_dropout,
activation_fn=self.activation_fn,
layer_norm_first=self.layer_norm_first,
encoder_layerdrop=self.encoder_layerdrop,
max_positions=self.max_positions,
)
## projections and dropouts
self.post_extract_proj = nn.Linear(self.extractor_embed, self.encoder_embed_dim)
self.dropout_input = nn.Dropout(self.dropout_input)
self.dropout_features = nn.Dropout(self.dropout_features)
self.layer_norm = torch.nn.LayerNorm(self.extractor_embed)
self.final_proj = nn.Linear(self.encoder_embed_dim, self.encoder_embed_dim)
# EMA module
self.average_top_k_layers = average_top_k_layers
self.layer_norm_target_layer = layer_norm_target_layer
self.instance_norm_target_layer = instance_norm_target_layer
self.instance_norm_targets = instance_norm_targets
self.layer_norm_targets = layer_norm_targets
self.batch_norm_target_layer = batch_norm_target_layer
self.group_norm_target_layer = group_norm_target_layer
self.ema_decay = ema_decay
self.ema_end_decay = ema_end_decay
self.ema_anneal_end_step = ema_anneal_end_step
self.ema_transformer_only = ema_transformer_only
self.ema_layers_only = ema_layers_only
self.min_target_var = min_target_var
self.min_pred_var = min_pred_var
self.ema = None
# Loss
self.loss_beta = loss_beta
self.loss_scale = loss_scale
# FP16 optimization
self.required_seq_len_multiple = required_seq_len_multiple
self.num_updates = 0
logging.info("Data2VecEncoder settings: {}".format(self.__dict__))
def make_ema_teacher(self):
"""Make ema teacher."""
skip_keys = set()
if self.ema_layers_only:
self.ema_transformer_only = True
for k, _ in self.encoder.pos_conv.named_parameters():
skip_keys.add(f"pos_conv.{k}")
self.ema = EMAModule(
self.encoder if self.ema_transformer_only else self,
ema_decay=self.ema_decay,
ema_fp32=True,
skip_keys=skip_keys,
)
def set_num_updates(self, num_updates):
"""Set num updates.
Args:
num_updates: TODO.
"""
if self.ema is None and self.final_proj is not None:
logging.info("Making EMA Teacher")
self.make_ema_teacher()
elif self.training and self.ema is not None:
if self.ema_decay != self.ema_end_decay:
if num_updates >= self.ema_anneal_end_step:
decay = self.ema_end_decay
else:
decay = get_annealed_rate(
self.ema_decay,
self.ema_end_decay,
num_updates,
self.ema_anneal_end_step,
)
self.ema.set_decay(decay)
if self.ema.get_decay() < 1:
self.ema.step(self.encoder if self.ema_transformer_only else self)
self.num_updates = num_updates
def apply_mask(
self,
x,
padding_mask,
mask_indices=None,
mask_channel_indices=None,
):
"""Apply mask.
Args:
x: TODO.
padding_mask: TODO.
mask_indices: TODO.
mask_channel_indices: TODO.
"""
B, T, C = x.shape
if self.mask_channel_prob > 0 and self.mask_channel_before:
mask_channel_indices = compute_mask_indices(
(B, C),
None,
self.mask_channel_prob,
self.mask_channel_length,
self.mask_channel_selection,
self.mask_channel_other,
no_overlap=self.no_mask_channel_overlap,
min_space=self.mask_channel_min_space,
)
mask_channel_indices = (
torch.from_numpy(mask_channel_indices).to(x.device).unsqueeze(1).expand(-1, T, -1)
)
x[mask_channel_indices] = 0
if self.mask_prob > 0:
if mask_indices is None:
mask_indices = compute_mask_indices(
(B, T),
padding_mask,
self.mask_prob,
self.mask_length,
self.mask_selection,
self.mask_other,
min_masks=1,
no_overlap=self.no_mask_overlap,
min_space=self.mask_min_space,
require_same_masks=self.require_same_masks,
mask_dropout=self.mask_dropout,
)
mask_indices = torch.from_numpy(mask_indices).to(x.device)
x[mask_indices] = self.mask_emb
else:
mask_indices = None
if self.mask_channel_prob > 0 and not self.mask_channel_before:
if mask_channel_indices is None:
mask_channel_indices = compute_mask_indices(
(B, C),
None,
self.mask_channel_prob,
self.mask_channel_length,
self.mask_channel_selection,
self.mask_channel_other,
no_overlap=self.no_mask_channel_overlap,
min_space=self.mask_channel_min_space,
)
mask_channel_indices = (
torch.from_numpy(mask_channel_indices)
.to(x.device)
.unsqueeze(1)
.expand(-1, T, -1)
)
x[mask_channel_indices] = 0
return x, mask_indices
def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
"""
Computes the output length of the convolutional layers
"""
def _conv_out_length(input_length, kernel_size, stride):
"""Internal: conv out length.
Args:
input_length: TODO.
kernel_size: Size/dimension parameter.
stride: TODO.
"""
return torch.floor((input_length - kernel_size).to(torch.float32) / stride + 1)
conv_cfg_list = eval(self.conv_feature_layers)
for i in range(len(conv_cfg_list)):
input_lengths = _conv_out_length(
input_lengths, conv_cfg_list[i][1], conv_cfg_list[i][2]
)
return input_lengths.to(torch.long)
def forward(
self,
xs_pad,
ilens=None,
mask=False,
features_only=True,
layer=None,
mask_indices=None,
mask_channel_indices=None,
padding_count=None,
):
# create padding_mask by ilens
"""Forward pass for training.
Args:
xs_pad: TODO.
ilens: TODO.
mask: TODO.
features_only: TODO.
layer: TODO.
mask_indices: TODO.
mask_channel_indices: TODO.
padding_count: TODO.
"""
if ilens is not None:
padding_mask = make_pad_mask(lengths=ilens).to(xs_pad.device)
else:
padding_mask = None
features = xs_pad
if self.feature_grad_mult > 0:
features = self.feature_extractor(features)
if self.feature_grad_mult != 1.0:
features = GradMultiply.apply(features, self.feature_grad_mult)
else:
with torch.no_grad():
features = self.feature_extractor(features)
features = features.transpose(1, 2)
features = self.layer_norm(features)
orig_padding_mask = padding_mask
if padding_mask is not None:
input_lengths = (1 - padding_mask.long()).sum(-1)
# apply conv formula to get real output_lengths
output_lengths = self._get_feat_extract_output_lengths(input_lengths)
padding_mask = torch.zeros(
features.shape[:2], dtype=features.dtype, device=features.device
)
# these two operations makes sure that all values
# before the output lengths indices are attended to
padding_mask[
(
torch.arange(padding_mask.shape[0], device=padding_mask.device),
output_lengths - 1,
)
] = 1
padding_mask = (1 - padding_mask.flip([-1]).cumsum(-1).flip([-1])).bool()
else:
padding_mask = None
if self.post_extract_proj is not None:
features = self.post_extract_proj(features)
pre_encoder_features = None
if self.ema_transformer_only:
pre_encoder_features = features.clone()
features = self.dropout_input(features)
if mask:
x, mask_indices = self.apply_mask(
features,
padding_mask,
mask_indices=mask_indices,
mask_channel_indices=mask_channel_indices,
)
else:
x = features
mask_indices = None
x, layer_results = self.encoder(
x,
padding_mask=padding_mask,
layer=layer,
)
if features_only:
encoder_out_lens = (1 - padding_mask.long()).sum(1)
return x, encoder_out_lens, None
result = {
"losses": {},
"padding_mask": padding_mask,
"x": x,
}
with torch.no_grad():
self.ema.model.eval()
if self.ema_transformer_only:
y, layer_results = self.ema.model.extract_features(
pre_encoder_features,
padding_mask=padding_mask,
min_layer=self.encoder_layers - self.average_top_k_layers,
)
y = {
"x": y,
"padding_mask": padding_mask,
"layer_results": layer_results,
}
else:
y = self.ema.model.extract_features(
source=xs_pad,
padding_mask=orig_padding_mask,
mask=False,
)
target_layer_results = [l[2] for l in y["layer_results"]]
permuted = False
if self.instance_norm_target_layer or self.batch_norm_target_layer:
target_layer_results = [
tl.permute(1, 2, 0) for tl in target_layer_results # TBC -> BCT
]
permuted = True
if self.batch_norm_target_layer:
target_layer_results = [
F.batch_norm(tl.float(), running_mean=None, running_var=None, training=True)
for tl in target_layer_results
]
if self.instance_norm_target_layer:
target_layer_results = [F.instance_norm(tl.float()) for tl in target_layer_results]
if permuted:
target_layer_results = [
tl.transpose(1, 2) for tl in target_layer_results # BCT -> BTC
]
if self.group_norm_target_layer:
target_layer_results = [
F.layer_norm(tl.float(), tl.shape[-2:]) for tl in target_layer_results
]
if self.layer_norm_target_layer:
target_layer_results = [
F.layer_norm(tl.float(), tl.shape[-1:]) for tl in target_layer_results
]
y = sum(target_layer_results) / len(target_layer_results)
if self.layer_norm_targets:
y = F.layer_norm(y.float(), y.shape[-1:])
if self.instance_norm_targets:
y = F.instance_norm(y.float().transpose(1, 2)).transpose(1, 2)
if not permuted:
y = y.transpose(0, 1)
y = y[mask_indices]
x = x[mask_indices]
x = self.final_proj(x)
sz = x.size(-1)
if self.loss_beta == 0:
loss = F.mse_loss(x.float(), y.float(), reduction="none").sum(dim=-1)
else:
loss = F.smooth_l1_loss(
x.float(), y.float(), reduction="none", beta=self.loss_beta
).sum(dim=-1)
if self.loss_scale is not None:
scale = self.loss_scale
else:
scale = 1 / math.sqrt(sz)
result["losses"]["regression"] = loss.sum() * scale
if "sample_size" not in result:
result["sample_size"] = loss.numel()
with torch.no_grad():
result["target_var"] = self.compute_var(y)
result["pred_var"] = self.compute_var(x.float())
if self.num_updates > 5000 and result["target_var"] < self.min_target_var:
logging.error(
f"target var is {result['target_var'].item()} < {self.min_target_var}, exiting"
)
raise Exception(
f"target var is {result['target_var'].item()} < {self.min_target_var}, exiting"
)
if self.num_updates > 5000 and result["pred_var"] < self.min_pred_var:
logging.error(f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting")
raise Exception(
f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting"
)
if self.ema is not None:
result["ema_decay"] = self.ema.get_decay() * 1000
return result
@staticmethod
def compute_var(y):
"""Compute var.
Args:
y: TODO.
"""
y = y.view(-1, y.size(-1))
if dist.is_initialized():
zc = torch.tensor(y.size(0)).cuda()
zs = y.sum(dim=0)
zss = (y**2).sum(dim=0)
dist.all_reduce(zc)
dist.all_reduce(zs)
dist.all_reduce(zss)
var = zss / (zc - 1) - (zs**2) / (zc * (zc - 1))
return torch.sqrt(var + 1e-6).mean()
else:
return torch.sqrt(y.var(dim=0) + 1e-6).mean()
def extract_features(self, xs_pad, ilens, mask=False, layer=None):
"""Extract features.
Args:
xs_pad: TODO.
ilens: TODO.
mask: TODO.
layer: TODO.
"""
res = self.forward(
xs_pad,
ilens,
mask=mask,
features_only=True,
layer=layer,
)
return res
def remove_pretraining_modules(self, last_layer=None):
"""Remove pretraining modules.
Args:
last_layer: TODO.
"""
self.final_proj = None
self.ema = None
if last_layer is not None:
self.encoder.layers = nn.ModuleList(
l for i, l in enumerate(self.encoder.layers) if i <= last_layer
)
def output_size(self) -> int:
"""Output size."""
return self.encoder_embed_dim
funasr/models/data2vec/data_utils.py
+149
View File
@@ -0,0 +1,149 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from typing import Optional, Tuple
import numpy as np
import torch
def compute_mask_indices(
shape: Tuple[int, int],
padding_mask: Optional[torch.Tensor],
mask_prob: float,
mask_length: int,
mask_type: str = "static",
mask_other: float = 0.0,
min_masks: int = 0,
no_overlap: bool = False,
min_space: int = 0,
require_same_masks: bool = True,
mask_dropout: float = 0.0,
) -> np.ndarray:
"""
Computes random mask spans for a given shape
Args:
shape: the the shape for which to compute masks.
should be of size 2 where first element is batch size and 2nd is timesteps
padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
however due to overlaps, the actual number will be smaller (unless no_overlap is True)
mask_type: how to compute mask lengths
static = fixed size
uniform = sample from uniform distribution [mask_other, mask_length*2]
normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
poisson = sample from possion distribution with lambda = mask length
min_masks: minimum number of masked spans
no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample
mask_dropout: randomly dropout this percentage of masks in each example
"""
bsz, all_sz = shape
mask = np.full((bsz, all_sz), False)
all_num_mask = int(
# add a random number for probabilistic rounding
mask_prob * all_sz / float(mask_length)
+ np.random.rand()
)
all_num_mask = max(min_masks, all_num_mask)
mask_idcs = []
for i in range(bsz):
if padding_mask is not None:
sz = all_sz - padding_mask[i].long().sum().item()
num_mask = int(
# add a random number for probabilistic rounding
mask_prob * sz / float(mask_length)
+ np.random.rand()
)
num_mask = max(min_masks, num_mask)
else:
sz = all_sz
num_mask = all_num_mask
if mask_type == "static":
lengths = np.full(num_mask, mask_length)
elif mask_type == "uniform":
lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
elif mask_type == "normal":
lengths = np.random.normal(mask_length, mask_other, size=num_mask)
lengths = [max(1, int(round(x))) for x in lengths]
elif mask_type == "poisson":
lengths = np.random.poisson(mask_length, size=num_mask)
lengths = [int(round(x)) for x in lengths]
else:
raise Exception("unknown mask selection " + mask_type)
if sum(lengths) == 0:
lengths[0] = min(mask_length, sz - 1)
if no_overlap:
mask_idc = []
def arrange(s, e, length, keep_length):
"""Arrange.
Args:
s: TODO.
e: TODO.
length: TODO.
keep_length: TODO.
"""
span_start = np.random.randint(s, e - length)
mask_idc.extend(span_start + i for i in range(length))
new_parts = []
if span_start - s - min_space >= keep_length:
new_parts.append((s, span_start - min_space + 1))
if e - span_start - length - min_space > keep_length:
new_parts.append((span_start + length + min_space, e))
return new_parts
parts = [(0, sz)]
min_length = min(lengths)
for length in sorted(lengths, reverse=True):
lens = np.fromiter(
(e - s if e - s >= length + min_space else 0 for s, e in parts),
np.int32,
)
l_sum = np.sum(lens)
if l_sum == 0:
break
probs = lens / np.sum(lens)
c = np.random.choice(len(parts), p=probs)
s, e = parts.pop(c)
parts.extend(arrange(s, e, length, min_length))
mask_idc = np.asarray(mask_idc)
else:
min_len = min(lengths)
if sz - min_len <= num_mask:
min_len = sz - num_mask - 1
mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
mask_idc = np.asarray(
[mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])]
)
mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
min_len = min([len(m) for m in mask_idcs])
for i, mask_idc in enumerate(mask_idcs):
if len(mask_idc) > min_len and require_same_masks:
mask_idc = np.random.choice(mask_idc, min_len, replace=False)
if mask_dropout > 0:
num_holes = np.rint(len(mask_idc) * mask_dropout).astype(int)
mask_idc = np.random.choice(mask_idc, len(mask_idc) - num_holes, replace=False)
mask[i, mask_idc] = True
return mask
funasr/models/data2vec/ema_module.py
+146
View File
@@ -0,0 +1,146 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""
Used for EMA tracking a given pytorch module. The user is responsible for calling step()
and setting the appropriate decay
"""
import copy
import logging
import torch
class EMAModule:
"""Exponential Moving Average of Fairseq Models"""
def __init__(self, model, ema_decay=0.9999, ema_fp32=False, device=None, skip_keys=None):
"""
@param model model to initialize the EMA with
@param config EMAConfig object with configuration like
ema_decay, ema_update_freq, ema_fp32
@param device If provided, copy EMA to this device (e.g. gpu).
Otherwise EMA is in the same device as the model.
"""
self.decay = ema_decay
self.ema_fp32 = ema_fp32
self.model = copy.deepcopy(model)
self.model.requires_grad_(False)
self.skip_keys = skip_keys or set()
self.fp32_params = {}
if device is not None:
logging.info(f"Copying EMA model to device {device}")
self.model = self.model.to(device=device)
if self.ema_fp32:
self.build_fp32_params()
self.update_freq_counter = 0
def build_fp32_params(self, state_dict=None):
"""
Store a copy of the EMA params in fp32.
If state dict is passed, the EMA params is copied from
the provided state dict. Otherwise, it is copied from the
current EMA model parameters.
"""
if not self.ema_fp32:
raise RuntimeError(
"build_fp32_params should not be called if ema_fp32=False. "
"Use ema_fp32=True if this is really intended."
)
if state_dict is None:
state_dict = self.model.state_dict()
def _to_float(t):
"""Internal: to float.
Args:
t: TODO.
"""
return t.float() if torch.is_floating_point(t) else t
for param_key in state_dict:
if param_key in self.fp32_params:
self.fp32_params[param_key].copy_(state_dict[param_key])
else:
self.fp32_params[param_key] = _to_float(state_dict[param_key])
def restore(self, state_dict, build_fp32_params=False):
"""Load data from a model spec into EMA model"""
self.model.load_state_dict(state_dict, strict=False)
if build_fp32_params:
self.build_fp32_params(state_dict)
def set_decay(self, decay):
"""Set decay.
Args:
decay: TODO.
"""
self.decay = decay
def get_decay(self):
"""Get decay."""
return self.decay
def _step_internal(self, new_model):
"""One update of the EMA model based on new model weights"""
decay = self.decay
ema_state_dict = {}
ema_params = self.fp32_params if self.ema_fp32 else self.model.state_dict()
for key, param in new_model.state_dict().items():
if isinstance(param, dict):
continue
try:
ema_param = ema_params[key]
except KeyError:
ema_param = param.float().clone() if param.ndim == 1 else copy.deepcopy(param)
if param.shape != ema_param.shape:
raise ValueError(
"incompatible tensor shapes between model param and ema param"
+ "{} vs. {}".format(param.shape, ema_param.shape)
)
if "version" in key:
# Do not decay a model.version pytorch param
continue
if key in self.skip_keys or (
"num_batches_tracked" in key and ema_param.dtype == torch.int64
):
ema_param = param.to(dtype=ema_param.dtype).clone()
ema_params[key].copy_(ema_param)
else:
ema_param.mul_(decay)
ema_param.add_(param.to(dtype=ema_param.dtype), alpha=1 - decay)
ema_state_dict[key] = ema_param
self.restore(ema_state_dict, build_fp32_params=False)
def step(self, new_model):
"""Step.
Args:
new_model: New Model instance.
"""
self._step_internal(new_model)
def reverse(self, model):
"""
Load the model parameters from EMA model.
Useful for inference or fine-tuning from the EMA model.
"""
d = self.model.state_dict()
if "_ema" in d:
del d["_ema"]
model.load_state_dict(d, strict=False)
return model
funasr/models/data2vec/grad_multiply.py
+31
View File
@@ -0,0 +1,31 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch
class GradMultiply(torch.autograd.Function):
@staticmethod
def forward(ctx, x, scale):
"""Forward pass for training.
Args:
ctx: TODO.
x: TODO.
scale: TODO.
"""
ctx.scale = scale
res = x.new(x)
return res
@staticmethod
def backward(ctx, grad):
"""Backward.
Args:
ctx: TODO.
grad: TODO.
"""
return grad * ctx.scale, None
funasr/models/data2vec/multihead_attention.py
+692
View File
@@ -0,0 +1,692 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import logging
import math
from typing import Dict, List, Optional, Tuple
import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.nn import Parameter
from funasr.models.data2vec.quant_noise import quant_noise
class FairseqDropout(nn.Module):
def __init__(self, p, module_name=None):
"""Initialize FairseqDropout.
Args:
p: TODO.
module_name: TODO.
"""
super().__init__()
self.p = p
self.module_name = module_name
self.apply_during_inference = False
def forward(self, x, inplace: bool = False):
"""Forward pass for training.
Args:
x: TODO.
inplace: TODO.
"""
if self.p > 0 and (self.training or self.apply_during_inference):
return F.dropout(x, p=self.p, training=True, inplace=inplace)
else:
return x
def make_generation_fast_(
self,
name: str,
retain_dropout: bool = False,
retain_dropout_modules: Optional[List[str]] = None,
**kwargs,
):
"""Make generation fast .
Args:
name: TODO.
retain_dropout: TODO.
retain_dropout_modules: TODO.
**kwargs: Additional keyword arguments.
"""
if retain_dropout:
if retain_dropout_modules is not None and self.module_name is None:
logging.warning(
"Cannot enable dropout during inference for module {} "
"because module_name was not set".format(name)
)
elif (
retain_dropout_modules is None # if None, apply to all modules
or self.module_name in retain_dropout_modules
):
logging.info("Enabling dropout during inference for module: {}".format(name))
self.apply_during_inference = True
else:
logging.info("Disabling dropout for module: {}".format(name))
class MultiheadAttention(nn.Module):
"""Multi-headed attention.
See "Attention Is All You Need" for more details.
"""
def __init__(
self,
embed_dim,
num_heads,
kdim=None,
vdim=None,
dropout=0.0,
bias=True,
add_bias_kv=False,
add_zero_attn=False,
self_attention=False,
encoder_decoder_attention=False,
q_noise=0.0,
qn_block_size=8,
):
"""Initialize MultiheadAttention.
Args:
embed_dim: Size/dimension parameter.
num_heads: TODO.
kdim: TODO.
vdim: TODO.
dropout: TODO.
bias: TODO.
add_bias_kv: TODO.
add_zero_attn: TODO.
self_attention: TODO.
encoder_decoder_attention: TODO.
q_noise: TODO.
qn_block_size: Size/dimension parameter.
"""
super().__init__()
self.embed_dim = embed_dim
self.kdim = kdim if kdim is not None else embed_dim
self.vdim = vdim if vdim is not None else embed_dim
self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
self.num_heads = num_heads
self.dropout_module = FairseqDropout(dropout, module_name=self.__class__.__name__)
self.head_dim = embed_dim // num_heads
assert (
self.head_dim * num_heads == self.embed_dim
), "embed_dim must be divisible by num_heads"
self.scaling = self.head_dim**-0.5
self.self_attention = self_attention
self.encoder_decoder_attention = encoder_decoder_attention
assert not self.self_attention or self.qkv_same_dim, (
"Self-attention requires query, key and " "value to be of the same size"
)
self.k_proj = quant_noise(
nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size
)
self.v_proj = quant_noise(
nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size
)
self.q_proj = quant_noise(
nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
)
self.out_proj = quant_noise(
nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
)
if add_bias_kv:
self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
else:
self.bias_k = self.bias_v = None
self.add_zero_attn = add_zero_attn
self.reset_parameters()
self.onnx_trace = False
self.skip_embed_dim_check = False
def prepare_for_onnx_export_(self):
"""Prepare for onnx export ."""
self.onnx_trace = True
def reset_parameters(self):
"""Reset parameters."""
if self.qkv_same_dim:
# Empirically observed the convergence to be much better with
# the scaled initialization
nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
else:
nn.init.xavier_uniform_(self.k_proj.weight)
nn.init.xavier_uniform_(self.v_proj.weight)
nn.init.xavier_uniform_(self.q_proj.weight)
nn.init.xavier_uniform_(self.out_proj.weight)
if self.out_proj.bias is not None:
nn.init.constant_(self.out_proj.bias, 0.0)
if self.bias_k is not None:
nn.init.xavier_normal_(self.bias_k)
if self.bias_v is not None:
nn.init.xavier_normal_(self.bias_v)
def _get_reserve_head_index(self, num_heads_to_keep: int):
"""Internal: get reserve head index.
Args:
num_heads_to_keep: TODO.
"""
k_proj_heads_norm = []
q_proj_heads_norm = []
v_proj_heads_norm = []
for i in range(self.num_heads):
start_idx = i * self.head_dim
end_idx = (i + 1) * self.head_dim
k_proj_heads_norm.append(
torch.sum(torch.abs(self.k_proj.weight[start_idx:end_idx,])).tolist()
+ torch.sum(torch.abs(self.k_proj.bias[start_idx:end_idx])).tolist()
)
q_proj_heads_norm.append(
torch.sum(torch.abs(self.q_proj.weight[start_idx:end_idx,])).tolist()
+ torch.sum(torch.abs(self.q_proj.bias[start_idx:end_idx])).tolist()
)
v_proj_heads_norm.append(
torch.sum(torch.abs(self.v_proj.weight[start_idx:end_idx,])).tolist()
+ torch.sum(torch.abs(self.v_proj.bias[start_idx:end_idx])).tolist()
)
heads_norm = []
for i in range(self.num_heads):
heads_norm.append(k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i])
sorted_head_index = sorted(range(self.num_heads), key=lambda k: heads_norm[k], reverse=True)
reserve_head_index = []
for i in range(num_heads_to_keep):
start = sorted_head_index[i] * self.head_dim
end = (sorted_head_index[i] + 1) * self.head_dim
reserve_head_index.append((start, end))
return reserve_head_index
def _adaptive_prune_heads(self, reserve_head_index: List[Tuple[int, int]]):
"""Internal: adaptive prune heads.
Args:
reserve_head_index: TODO.
"""
new_q_weight = []
new_q_bias = []
new_k_weight = []
new_k_bias = []
new_v_weight = []
new_v_bias = []
new_out_proj_weight = []
for ele in reserve_head_index:
start_idx, end_idx = ele
new_q_weight.append(self.q_proj.weight[start_idx:end_idx,])
new_q_bias.append(self.q_proj.bias[start_idx:end_idx])
new_k_weight.append(self.k_proj.weight[start_idx:end_idx,])
new_k_bias.append(self.k_proj.bias[start_idx:end_idx])
new_v_weight.append(self.v_proj.weight[start_idx:end_idx,])
new_v_bias.append(self.v_proj.bias[start_idx:end_idx])
new_out_proj_weight.append(self.out_proj.weight[:, start_idx:end_idx])
new_q_weight = torch.cat(new_q_weight).detach()
new_k_weight = torch.cat(new_k_weight).detach()
new_v_weight = torch.cat(new_v_weight).detach()
new_out_proj_weight = torch.cat(new_out_proj_weight, dim=-1).detach()
new_q_weight.requires_grad = True
new_k_weight.requires_grad = True
new_v_weight.requires_grad = True
new_out_proj_weight.requires_grad = True
new_q_bias = torch.cat(new_q_bias).detach()
new_q_bias.requires_grad = True
new_k_bias = torch.cat(new_k_bias).detach()
new_k_bias.requires_grad = True
new_v_bias = torch.cat(new_v_bias).detach()
new_v_bias.requires_grad = True
self.q_proj.weight = torch.nn.Parameter(new_q_weight)
self.q_proj.bias = torch.nn.Parameter(new_q_bias)
self.k_proj.weight = torch.nn.Parameter(new_k_weight)
self.k_proj.bias = torch.nn.Parameter(new_k_bias)
self.v_proj.weight = torch.nn.Parameter(new_v_weight)
self.v_proj.bias = torch.nn.Parameter(new_v_bias)
self.out_proj.weight = torch.nn.Parameter(new_out_proj_weight)
self.num_heads = len(reserve_head_index)
self.embed_dim = self.head_dim * self.num_heads
self.q_proj.out_features = self.embed_dim
self.k_proj.out_features = self.embed_dim
self.v_proj.out_features = self.embed_dim
def _set_skip_embed_dim_check(self):
"""Internal: set skip embed dim check."""
self.skip_embed_dim_check = True
def forward(
self,
query,
key: Optional[Tensor],
value: Optional[Tensor],
key_padding_mask: Optional[Tensor] = None,
incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
need_weights: bool = True,
static_kv: bool = False,
attn_mask: Optional[Tensor] = None,
before_softmax: bool = False,
need_head_weights: bool = False,
) -> Tuple[Tensor, Optional[Tensor]]:
"""Input shape: Time x Batch x Channel
Args:
key_padding_mask (ByteTensor, optional): mask to exclude
keys that are pads, of shape `(batch, src_len)`, where
padding elements are indicated by 1s.
need_weights (bool, optional): return the attention weights,
averaged over heads (default: False).
attn_mask (ByteTensor, optional): typically used to
implement causal attention, where the mask prevents the
attention from looking forward in time (default: None).
before_softmax (bool, optional): return the raw attention
weights and values before the attention softmax.
need_head_weights (bool, optional): return the attention
weights for each head. Implies *need_weights*. Default:
return the average attention weights over all heads.
"""
if need_head_weights:
need_weights = True
is_tpu = query.device.type == "xla"
tgt_len, bsz, embed_dim = query.size()
src_len = tgt_len
if not self.skip_embed_dim_check:
assert embed_dim == self.embed_dim, f"query dim {embed_dim} != {self.embed_dim}"
assert list(query.size()) == [tgt_len, bsz, embed_dim]
if key is not None:
src_len, key_bsz, _ = key.size()
if not torch.jit.is_scripting():
assert key_bsz == bsz
assert value is not None
assert src_len, bsz == value.shape[:2]
if (
not self.onnx_trace
and not is_tpu # don't use PyTorch version on TPUs
and incremental_state is None
and not static_kv
# A workaround for quantization to work. Otherwise JIT compilation
# treats bias in linear module as method.
and not torch.jit.is_scripting()
# The Multihead attention implemented in pytorch forces strong dimension check
# for input embedding dimention and K,Q,V projection dimension.
# Since pruning will break the dimension check and it is not easy to modify the pytorch API,
# it is preferred to bypass the pytorch MHA when we need to skip embed_dim_check
and not self.skip_embed_dim_check
):
assert key is not None and value is not None
return F.multi_head_attention_forward(
query,
key,
value,
self.embed_dim,
self.num_heads,
torch.empty([0]),
torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),
self.bias_k,
self.bias_v,
self.add_zero_attn,
self.dropout_module.p,
self.out_proj.weight,
self.out_proj.bias,
self.training or self.dropout_module.apply_during_inference,
key_padding_mask,
need_weights,
attn_mask,
use_separate_proj_weight=True,
q_proj_weight=self.q_proj.weight,
k_proj_weight=self.k_proj.weight,
v_proj_weight=self.v_proj.weight,
)
if incremental_state is not None:
saved_state = self._get_input_buffer(incremental_state)
if saved_state is not None and "prev_key" in saved_state:
# previous time steps are cached - no need to recompute
# key and value if they are static
if static_kv:
assert self.encoder_decoder_attention and not self.self_attention
key = value = None
else:
saved_state = None
if self.self_attention:
q = self.q_proj(query)
k = self.k_proj(query)
v = self.v_proj(query)
elif self.encoder_decoder_attention:
# encoder-decoder attention
q = self.q_proj(query)
if key is None:
assert value is None
k = v = None
else:
k = self.k_proj(key)
v = self.v_proj(key)
else:
assert key is not None and value is not None
q = self.q_proj(query)
k = self.k_proj(key)
v = self.v_proj(value)
q *= self.scaling
if self.bias_k is not None:
assert self.bias_v is not None
k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
if attn_mask is not None:
attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
if key_padding_mask is not None:
key_padding_mask = torch.cat(
[
key_padding_mask,
key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
],
dim=1,
)
q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
if k is not None:
k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
if v is not None:
v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
if saved_state is not None:
# saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
if "prev_key" in saved_state:
_prev_key = saved_state["prev_key"]
assert _prev_key is not None
prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
if static_kv:
k = prev_key
else:
assert k is not None
k = torch.cat([prev_key, k], dim=1)
src_len = k.size(1)
if "prev_value" in saved_state:
_prev_value = saved_state["prev_value"]
assert _prev_value is not None
prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
if static_kv:
v = prev_value
else:
assert v is not None
v = torch.cat([prev_value, v], dim=1)
prev_key_padding_mask: Optional[Tensor] = None
if "prev_key_padding_mask" in saved_state:
prev_key_padding_mask = saved_state["prev_key_padding_mask"]
assert k is not None and v is not None
key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
key_padding_mask=key_padding_mask,
prev_key_padding_mask=prev_key_padding_mask,
batch_size=bsz,
src_len=k.size(1),
static_kv=static_kv,
)
saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
saved_state["prev_key_padding_mask"] = key_padding_mask
# In this branch incremental_state is never None
assert incremental_state is not None
incremental_state = self._set_input_buffer(incremental_state, saved_state)
assert k is not None
assert k.size(1) == src_len
# This is part of a workaround to get around fork/join parallelism
# not supporting Optional types.
if key_padding_mask is not None and key_padding_mask.dim() == 0:
key_padding_mask = None
if key_padding_mask is not None:
assert key_padding_mask.size(0) == bsz
assert key_padding_mask.size(1) == src_len
if self.add_zero_attn:
assert v is not None
src_len += 1
k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
if attn_mask is not None:
attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
if key_padding_mask is not None:
key_padding_mask = torch.cat(
[
key_padding_mask,
torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask),
],
dim=1,
)
attn_weights = torch.bmm(q, k.transpose(1, 2))
attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
if attn_mask is not None:
attn_mask = attn_mask.unsqueeze(0)
if self.onnx_trace:
attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1)
attn_weights += attn_mask
if key_padding_mask is not None:
# don't attend to padding symbols
attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
if not is_tpu:
attn_weights = attn_weights.masked_fill(
key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
float("-inf"),
)
else:
attn_weights = attn_weights.transpose(0, 2)
attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf"))
attn_weights = attn_weights.transpose(0, 2)
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
if before_softmax:
return attn_weights, v
attn_weights_float = F.softmax(attn_weights, dim=-1, dtype=torch.float32)
attn_weights = attn_weights_float.type_as(attn_weights)
attn_probs = self.dropout_module(attn_weights)
assert v is not None
attn = torch.bmm(attn_probs, v)
assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
if self.onnx_trace and attn.size(1) == 1:
# when ONNX tracing a single decoder step (sequence length == 1)
# the transpose is a no-op copy before view, thus unnecessary
attn = attn.contiguous().view(tgt_len, bsz, self.embed_dim)
else:
attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim)
attn = self.out_proj(attn)
attn_weights: Optional[Tensor] = None
if need_weights:
attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose(
1, 0
)
if not need_head_weights:
# average attention weights over heads
attn_weights = attn_weights.mean(dim=0)
return attn, attn_weights
@staticmethod
def _append_prev_key_padding_mask(
key_padding_mask: Optional[Tensor],
prev_key_padding_mask: Optional[Tensor],
batch_size: int,
src_len: int,
static_kv: bool,
) -> Optional[Tensor]:
# saved key padding masks have shape (bsz, seq_len)
"""Internal: append prev key padding mask.
Args:
key_padding_mask: TODO.
prev_key_padding_mask: TODO.
batch_size: Number of samples per batch.
src_len: TODO.
static_kv: TODO.
"""
if prev_key_padding_mask is not None and static_kv:
new_key_padding_mask = prev_key_padding_mask
elif prev_key_padding_mask is not None and key_padding_mask is not None:
new_key_padding_mask = torch.cat(
[prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
)
# During incremental decoding, as the padding token enters and
# leaves the frame, there will be a time when prev or current
# is None
elif prev_key_padding_mask is not None:
if src_len > prev_key_padding_mask.size(1):
filler = torch.zeros(
(batch_size, src_len - prev_key_padding_mask.size(1)),
device=prev_key_padding_mask.device,
)
new_key_padding_mask = torch.cat(
[prev_key_padding_mask.float(), filler.float()], dim=1
)
else:
new_key_padding_mask = prev_key_padding_mask.float()
elif key_padding_mask is not None:
if src_len > key_padding_mask.size(1):
filler = torch.zeros(
(batch_size, src_len - key_padding_mask.size(1)),
device=key_padding_mask.device,
)
new_key_padding_mask = torch.cat([filler.float(), key_padding_mask.float()], dim=1)
else:
new_key_padding_mask = key_padding_mask.float()
else:
new_key_padding_mask = prev_key_padding_mask
return new_key_padding_mask
@torch.jit.export
def reorder_incremental_state(
self,
incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
new_order: Tensor,
):
"""Reorder buffered internal state (for incremental generation)."""
input_buffer = self._get_input_buffer(incremental_state)
if input_buffer is not None:
for k in input_buffer.keys():
input_buffer_k = input_buffer[k]
if input_buffer_k is not None:
if self.encoder_decoder_attention and input_buffer_k.size(0) == new_order.size(
0
):
break
input_buffer[k] = input_buffer_k.index_select(0, new_order)
incremental_state = self._set_input_buffer(incremental_state, input_buffer)
return incremental_state
def _get_input_buffer(
self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
) -> Dict[str, Optional[Tensor]]:
"""Internal: get input buffer.
Args:
incremental_state: TODO.
"""
result = self.get_incremental_state(incremental_state, "attn_state")
if result is not None:
return result
else:
empty_result: Dict[str, Optional[Tensor]] = {}
return empty_result
def _set_input_buffer(
self,
incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
buffer: Dict[str, Optional[Tensor]],
):
"""Internal: set input buffer.
Args:
incremental_state: TODO.
buffer: TODO.
"""
return self.set_incremental_state(incremental_state, "attn_state", buffer)
def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int):
"""Apply sparse mask.
Args:
attn_weights: TODO.
tgt_len: TODO.
src_len: TODO.
bsz: TODO.
"""
return attn_weights
def upgrade_state_dict_named(self, state_dict, name):
"""Upgrade state dict named.
Args:
state_dict: TODO.
name: TODO.
"""
prefix = name + "." if name != "" else ""
items_to_add = {}
keys_to_remove = []
for k in state_dict.keys():
if k.endswith(prefix + "in_proj_weight"):
# in_proj_weight used to be q + k + v with same dimensions
dim = int(state_dict[k].shape[0] / 3)
items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim]
items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim]
items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :]
keys_to_remove.append(k)
k_bias = prefix + "in_proj_bias"
if k_bias in state_dict.keys():
dim = int(state_dict[k].shape[0] / 3)
items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]
items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][dim : 2 * dim]
items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]
keys_to_remove.append(prefix + "in_proj_bias")
for k in keys_to_remove:
del state_dict[k]
for key, value in items_to_add.items():
state_dict[key] = value
funasr/models/data2vec/quant_noise.py
+107
View File
@@ -0,0 +1,107 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch
import torch.nn as nn
def quant_noise(module, p, block_size):
"""
Wraps modules and applies quantization noise to the weights for
subsequent quantization with Iterative Product Quantization as
described in "Training with Quantization Noise for Extreme Model Compression"
Args:
- module: nn.Module
- p: amount of Quantization Noise
- block_size: size of the blocks for subsequent quantization with iPQ
Remarks:
- Module weights must have the right sizes wrt the block size
- Only Linear, Embedding and Conv2d modules are supported for the moment
- For more detail on how to quantize by blocks with convolutional weights,
see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks"
- We implement the simplest form of noise here as stated in the paper
which consists in randomly dropping blocks
"""
# if no quantization noise, don't register hook
if p <= 0:
return module
# supported modules
assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
# test whether module.weight has the right sizes wrt block_size
is_conv = module.weight.ndim == 4
# 2D matrix
if not is_conv:
assert (
module.weight.size(1) % block_size == 0
), "Input features must be a multiple of block sizes"
# 4D matrix
else:
# 1x1 convolutions
if module.kernel_size == (1, 1):
assert (
module.in_channels % block_size == 0
), "Input channels must be a multiple of block sizes"
# regular convolutions
else:
k = module.kernel_size[0] * module.kernel_size[1]
assert k % block_size == 0, "Kernel size must be a multiple of block size"
def _forward_pre_hook(mod, input):
# no noise for evaluation
"""Internal: forward pre hook.
Args:
mod: TODO.
input: Input audio/text data.
"""
if mod.training:
if not is_conv:
# gather weight and sizes
weight = mod.weight
in_features = weight.size(1)
out_features = weight.size(0)
# split weight matrix into blocks and randomly drop selected blocks
mask = torch.zeros(in_features // block_size * out_features, device=weight.device)
mask.bernoulli_(p)
mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
else:
# gather weight and sizes
weight = mod.weight
in_channels = mod.in_channels
out_channels = mod.out_channels
# split weight matrix into blocks and randomly drop selected blocks
if mod.kernel_size == (1, 1):
mask = torch.zeros(
int(in_channels // block_size * out_channels),
device=weight.device,
)
mask.bernoulli_(p)
mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
else:
mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device)
mask.bernoulli_(p)
mask = (
mask.unsqueeze(2)
.unsqueeze(3)
.repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
)
# scale weights and apply mask
mask = mask.to(torch.bool) # x.bool() is not currently supported in TorchScript
s = 1 / (1 - p)
mod.weight.data = s * weight.masked_fill(mask, 0)
module.register_forward_pre_hook(_forward_pre_hook)
return module
funasr/models/data2vec/utils.py
+221
View File
@@ -0,0 +1,221 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from funasr.models.data2vec.multihead_attention import MultiheadAttention
class Fp32LayerNorm(nn.LayerNorm):
def __init__(self, *args, **kwargs):
"""Initialize Fp32LayerNorm.
Args:
*args: Variable positional arguments.
**kwargs: Additional keyword arguments.
"""
super().__init__(*args, **kwargs)
def forward(self, input):
"""Forward pass for training.
Args:
input: Input audio/text data.
"""
output = F.layer_norm(
input.float(),
self.normalized_shape,
self.weight.float() if self.weight is not None else None,
self.bias.float() if self.bias is not None else None,
self.eps,
)
return output.type_as(input)
class Fp32GroupNorm(nn.GroupNorm):
def __init__(self, *args, **kwargs):
"""Initialize Fp32GroupNorm.
Args:
*args: Variable positional arguments.
**kwargs: Additional keyword arguments.
"""
super().__init__(*args, **kwargs)
def forward(self, input):
"""Forward pass for training.
Args:
input: Input audio/text data.
"""
output = F.group_norm(
input.float(),
self.num_groups,
self.weight.float() if self.weight is not None else None,
self.bias.float() if self.bias is not None else None,
self.eps,
)
return output.type_as(input)
class TransposeLast(nn.Module):
def __init__(self, deconstruct_idx=None):
"""Initialize TransposeLast.
Args:
deconstruct_idx: TODO.
"""
super().__init__()
self.deconstruct_idx = deconstruct_idx
def forward(self, x):
"""Forward pass for training.
Args:
x: TODO.
"""
if self.deconstruct_idx is not None:
x = x[self.deconstruct_idx]
return x.transpose(-2, -1)
class SamePad(nn.Module):
def __init__(self, kernel_size, causal=False):
"""Initialize SamePad.
Args:
kernel_size: Size/dimension parameter.
causal: TODO.
"""
super().__init__()
if causal:
self.remove = kernel_size - 1
else:
self.remove = 1 if kernel_size % 2 == 0 else 0
def forward(self, x):
"""Forward pass for training.
Args:
x: TODO.
"""
if self.remove > 0:
x = x[:, :, : -self.remove]
return x
def pad_to_multiple(x, multiple, dim=-1, value=0):
# Inspired from https://github.com/lucidrains/local-attention/blob/master/local_attention/local_attention.py#L41
"""Pad to multiple.
Args:
x: TODO.
multiple: TODO.
dim: TODO.
value: TODO.
"""
if x is None:
return None, 0
tsz = x.size(dim)
m = tsz / multiple
remainder = math.ceil(m) * multiple - tsz
if m.is_integer():
return x, 0
pad_offset = (0,) * (-1 - dim) * 2
return F.pad(x, (*pad_offset, 0, remainder), value=value), remainder
def gelu_accurate(x):
"""Gelu accurate.
Args:
x: TODO.
"""
if not hasattr(gelu_accurate, "_a"):
gelu_accurate._a = math.sqrt(2 / math.pi)
return 0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3))))
def gelu(x: torch.Tensor) -> torch.Tensor:
"""Gelu.
Args:
x: TODO.
"""
return torch.nn.functional.gelu(x.float()).type_as(x)
def get_available_activation_fns():
"""Get available activation fns."""
return [
"relu",
"gelu",
"gelu_fast", # deprecated
"gelu_accurate",
"tanh",
"linear",
]
def get_activation_fn(activation: str):
"""Returns the activation function corresponding to `activation`"""
if activation == "relu":
return F.relu
elif activation == "gelu":
return gelu
elif activation == "gelu_accurate":
return gelu_accurate
elif activation == "tanh":
return torch.tanh
elif activation == "linear":
return lambda x: x
elif activation == "swish":
return torch.nn.SiLU
else:
raise RuntimeError("--activation-fn {} not supported".format(activation))
def init_bert_params(module):
"""
Initialize the weights specific to the BERT Model.
This overrides the default initializations depending on the specified arguments.
1. If normal_init_linear_weights is set then weights of linear
layer will be initialized using the normal distribution and
bais will be set to the specified value.
2. If normal_init_embed_weights is set then weights of embedding
layer will be initialized using the normal distribution.
3. If normal_init_proj_weights is set then weights of
in_project_weight for MultiHeadAttention initialized using
the normal distribution (to be validated).
"""
def normal_(data):
# with FSDP, module params will be on CUDA, so we cast them back to CPU
# so that the RNG is consistent with and without FSDP
"""Normal .
Args:
data: TODO.
"""
data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
if isinstance(module, nn.Linear):
normal_(module.weight.data)
if module.bias is not None:
module.bias.data.zero_()
if isinstance(module, nn.Embedding):
normal_(module.weight.data)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
if isinstance(module, MultiheadAttention):
normal_(module.q_proj.weight.data)
normal_(module.k_proj.weight.data)
normal_(module.v_proj.weight.data)
funasr/models/data2vec/wav2vec2.py
+497
View File
@@ -0,0 +1,497 @@
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import logging
import math
from typing import List, Tuple
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from funasr.models.data2vec import utils
from funasr.models.data2vec.multihead_attention import MultiheadAttention
class ConvFeatureExtractionModel(nn.Module):
def __init__(
self,
conv_layers: List[Tuple[int, int, int]],
dropout: float = 0.0,
mode: str = "default",
conv_bias: bool = False,
in_d: int = 1,
):
"""Initialize ConvFeatureExtractionModel.
Args:
conv_layers: TODO.
dropout: TODO.
mode: TODO.
conv_bias: TODO.
in_d: TODO.
"""
super().__init__()
assert mode in {"default", "layer_norm"}
def block(
n_in,
n_out,
k,
stride,
is_layer_norm=False,
is_group_norm=False,
conv_bias=False,
):
"""Block.
Args:
n_in: TODO.
n_out: TODO.
k: TODO.
stride: TODO.
is_layer_norm: Boolean flag for layer norm.
is_group_norm: Boolean flag for group norm.
conv_bias: TODO.
"""
def make_conv():
"""Make conv."""
conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
nn.init.kaiming_normal_(conv.weight)
return conv
assert (
is_layer_norm and is_group_norm
) == False, "layer norm and group norm are exclusive"
if is_layer_norm:
return nn.Sequential(
make_conv(),
nn.Dropout(p=dropout),
nn.Sequential(
utils.TransposeLast(),
utils.Fp32LayerNorm(dim, elementwise_affine=True),
utils.TransposeLast(),
),
nn.GELU(),
)
elif is_group_norm:
return nn.Sequential(
make_conv(),
nn.Dropout(p=dropout),
utils.Fp32GroupNorm(dim, dim, affine=True),
nn.GELU(),
)
else:
return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU())
self.conv_layers = nn.ModuleList()
for i, cl in enumerate(conv_layers):
assert len(cl) == 3, "invalid conv definition: " + str(cl)
(dim, k, stride) = cl
self.conv_layers.append(
block(
in_d,
dim,
k,
stride,
is_layer_norm=mode == "layer_norm",
is_group_norm=mode == "default" and i == 0,
conv_bias=conv_bias,
)
)
in_d = dim
def forward(self, x):
"""Forward pass for training.
Args:
x: TODO.
"""
if len(x.shape) == 2:
x = x.unsqueeze(1)
else:
x = x.transpose(1, 2)
for conv in self.conv_layers:
x = conv(x)
return x
def make_conv_pos(e, k, g):
"""Make conv pos.
Args:
e: TODO.
k: TODO.
g: TODO.
"""
pos_conv = nn.Conv1d(
e,
e,
kernel_size=k,
padding=k // 2,
groups=g,
)
dropout = 0
std = math.sqrt((4 * (1.0 - dropout)) / (k * e))
nn.init.normal_(pos_conv.weight, mean=0, std=std)
nn.init.constant_(pos_conv.bias, 0)
pos_conv = nn.utils.weight_norm(pos_conv, name="weight", dim=2)
pos_conv = nn.Sequential(pos_conv, utils.SamePad(k), nn.GELU())
return pos_conv
class TransformerEncoder(nn.Module):
def build_encoder_layer(self):
"""Build encoder layer."""
if self.layer_type == "transformer":
layer = TransformerSentenceEncoderLayer(
embedding_dim=self.embedding_dim,
ffn_embedding_dim=self.encoder_ffn_embed_dim,
num_attention_heads=self.encoder_attention_heads,
dropout=self.dropout,
attention_dropout=self.attention_dropout,
activation_dropout=self.activation_dropout,
activation_fn=self.activation_fn,
layer_norm_first=self.layer_norm_first,
)
else:
logging.error("Only transformer is supported for data2vec now")
return layer
def __init__(
self,
# position
dropout,
encoder_embed_dim,
required_seq_len_multiple,
pos_conv_depth,
conv_pos,
conv_pos_groups,
# transformer layers
layer_type,
encoder_layers,
encoder_ffn_embed_dim,
encoder_attention_heads,
attention_dropout,
activation_dropout,
activation_fn,
layer_norm_first,
encoder_layerdrop,
max_positions,
):
"""Initialize TransformerEncoder.
Args:
dropout: TODO.
encoder_embed_dim: Size/dimension parameter.
required_seq_len_multiple: TODO.
pos_conv_depth: TODO.
conv_pos: TODO.
conv_pos_groups: TODO.
layer_type: TODO.
encoder_layers: TODO.
encoder_ffn_embed_dim: Size/dimension parameter.
encoder_attention_heads: TODO.
attention_dropout: TODO.
activation_dropout: TODO.
activation_fn: TODO.
layer_norm_first: TODO.
encoder_layerdrop: TODO.
max_positions: TODO.
"""
super().__init__()
# position
self.dropout = dropout
self.embedding_dim = encoder_embed_dim
self.required_seq_len_multiple = required_seq_len_multiple
if pos_conv_depth > 1:
num_layers = pos_conv_depth
k = max(3, conv_pos // num_layers)
def make_conv_block(e, k, g, l):
"""Make conv block.
Args:
e: TODO.
k: TODO.
g: TODO.
l: TODO.
"""
return nn.Sequential(
*[
nn.Sequential(
nn.Conv1d(
e,
e,
kernel_size=k,
padding=k // 2,
groups=g,
),
utils.SamePad(k),
utils.TransposeLast(),
torch.nn.LayerNorm(e, elementwise_affine=False),
utils.TransposeLast(),
nn.GELU(),
)
for _ in range(l)
]
)
self.pos_conv = make_conv_block(self.embedding_dim, k, conv_pos_groups, num_layers)
else:
self.pos_conv = make_conv_pos(
self.embedding_dim,
conv_pos,
conv_pos_groups,
)
# transformer layers
self.layer_type = layer_type
self.encoder_ffn_embed_dim = encoder_ffn_embed_dim
self.encoder_attention_heads = encoder_attention_heads
self.attention_dropout = attention_dropout
self.activation_dropout = activation_dropout
self.activation_fn = activation_fn
self.layer_norm_first = layer_norm_first
self.layerdrop = encoder_layerdrop
self.max_positions = max_positions
self.layers = nn.ModuleList([self.build_encoder_layer() for _ in range(encoder_layers)])
self.layer_norm = torch.nn.LayerNorm(self.embedding_dim)
self.apply(utils.init_bert_params)
def forward(self, x, padding_mask=None, layer=None):
"""Forward pass for training.
Args:
x: TODO.
padding_mask: TODO.
layer: TODO.
"""
x, layer_results = self.extract_features(x, padding_mask, layer)
if self.layer_norm_first and layer is None:
x = self.layer_norm(x)
return x, layer_results
def extract_features(
self,
x,
padding_mask=None,
tgt_layer=None,
min_layer=0,
):
"""Extract features.
Args:
x: TODO.
padding_mask: TODO.
tgt_layer: TODO.
min_layer: TODO.
"""
if padding_mask is not None:
x[padding_mask] = 0
x_conv = self.pos_conv(x.transpose(1, 2))
x_conv = x_conv.transpose(1, 2)
x = x + x_conv
if not self.layer_norm_first:
x = self.layer_norm(x)
# pad to the sequence length dimension
x, pad_length = utils.pad_to_multiple(x, self.required_seq_len_multiple, dim=-2, value=0)
if pad_length > 0 and padding_mask is None:
padding_mask = x.new_zeros((x.size(0), x.size(1)), dtype=torch.bool)
padding_mask[:, -pad_length:] = True
else:
padding_mask, _ = utils.pad_to_multiple(
padding_mask, self.required_seq_len_multiple, dim=-1, value=True
)
x = F.dropout(x, p=self.dropout, training=self.training)
# B x T x C -> T x B x C
x = x.transpose(0, 1)
layer_results = []
r = None
for i, layer in enumerate(self.layers):
dropout_probability = np.random.random() if self.layerdrop > 0 else 1
if not self.training or (dropout_probability > self.layerdrop):
x, (z, lr) = layer(x, self_attn_padding_mask=padding_mask)
if i >= min_layer:
layer_results.append((x, z, lr))
if i == tgt_layer:
r = x
break
if r is not None:
x = r
# T x B x C -> B x T x C
x = x.transpose(0, 1)
# undo paddding
if pad_length > 0:
x = x[:, :-pad_length]
def undo_pad(a, b, c):
"""Undo pad.
Args:
a: TODO.
b: TODO.
c: TODO.
"""
return (
a[:-pad_length],
b[:-pad_length] if b is not None else b,
c[:-pad_length],
)
layer_results = [undo_pad(*u) for u in layer_results]
return x, layer_results
def max_positions(self):
"""Maximum output length supported by the encoder."""
return self.max_positions
def upgrade_state_dict_named(self, state_dict, name):
"""Upgrade a (possibly old) state dict for new versions of fairseq."""
return state_dict
class TransformerSentenceEncoderLayer(nn.Module):
"""
Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
models.
"""
def __init__(
self,
embedding_dim: int = 768,
ffn_embedding_dim: int = 3072,
num_attention_heads: int = 8,
dropout: float = 0.1,
attention_dropout: float = 0.1,
activation_dropout: float = 0.1,
activation_fn: str = "relu",
layer_norm_first: bool = False,
) -> None:
"""Initialize TransformerSentenceEncoderLayer.
Args:
embedding_dim: Size/dimension parameter.
ffn_embedding_dim: Size/dimension parameter.
num_attention_heads: TODO.
dropout: TODO.
attention_dropout: TODO.
activation_dropout: TODO.
activation_fn: TODO.
layer_norm_first: TODO.
"""
super().__init__()
# Initialize parameters
self.embedding_dim = embedding_dim
self.dropout = dropout
self.activation_dropout = activation_dropout
# Initialize blocks
self.activation_fn = utils.get_activation_fn(activation_fn)
self.self_attn = MultiheadAttention(
self.embedding_dim,
num_attention_heads,
dropout=attention_dropout,
self_attention=True,
)
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(self.activation_dropout)
self.dropout3 = nn.Dropout(dropout)
self.layer_norm_first = layer_norm_first
# layer norm associated with the self attention layer
self.self_attn_layer_norm = torch.nn.LayerNorm(self.embedding_dim)
self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
# layer norm associated with the position wise feed-forward NN
self.final_layer_norm = torch.nn.LayerNorm(self.embedding_dim)
def forward(
self,
x: torch.Tensor, # (T, B, C)
self_attn_mask: torch.Tensor = None,
self_attn_padding_mask: torch.Tensor = None,
):
"""
LayerNorm is applied either before or after the self-attention/ffn
modules similar to the original Transformer imlementation.
"""
residual = x
if self.layer_norm_first:
x = self.self_attn_layer_norm(x)
x, attn = self.self_attn(
query=x,
key=x,
value=x,
key_padding_mask=self_attn_padding_mask,
attn_mask=self_attn_mask,
need_weights=False,
)
x = self.dropout1(x)
x = residual + x
residual = x
x = self.final_layer_norm(x)
x = self.activation_fn(self.fc1(x))
x = self.dropout2(x)
x = self.fc2(x)
layer_result = x
x = self.dropout3(x)
x = residual + x
else:
x, attn = self.self_attn(
query=x,
key=x,
value=x,
key_padding_mask=self_attn_padding_mask,
need_weights=False,
)
x = self.dropout1(x)
x = residual + x
x = self.self_attn_layer_norm(x)
residual = x
x = self.activation_fn(self.fc1(x))
x = self.dropout2(x)
x = self.fc2(x)
layer_result = x
x = self.dropout3(x)
x = residual + x
x = self.final_layer_norm(x)
return x, (attn, layer_result)