Spaces:
Running
on
A10G
Running
on
A10G
from typing import Optional, Tuple | |
import torch | |
from wenet.ssl.bestrq.mask import compute_mask_indices | |
from wenet.utils.mask import make_pad_mask | |
class BestRQModel(torch.nn.Module): | |
def __init__( | |
self, | |
encoder: torch.nn.Module, | |
input_dim: int = 256, | |
embedding_dim: int = 256, | |
num_embeddings: int = 8192, | |
num_codebooks: int = 1, | |
dropout_rate: float = 0.1, | |
mask_prob: float = 0.01, | |
mask_length: int = 10, | |
min_masks: int = 2, | |
layer_norm_epsilon=1e-5, | |
) -> None: | |
super().__init__() | |
assert mask_prob > 0.0 | |
self.mask_prob = mask_prob | |
# NOTE: should filter audio less than mask_length | |
self.mask_length = mask_length | |
self.min_masks = min_masks | |
self.input_dropout = torch.nn.Dropout(dropout_rate) | |
# [embedding_dim, num_embeddings] | |
random_embedding_weight = torch.empty( | |
num_codebooks, embedding_dim, num_embeddings, requires_grad=False | |
) | |
self.embeddings = torch.nn.init.normal_(random_embedding_weight) | |
random_projection_weight = torch.empty( | |
input_dim, embedding_dim, requires_grad=False | |
) | |
self.projection = torch.nn.init.xavier_normal_(random_projection_weight) | |
mask_emb_weight = torch.Tensor(input_dim) | |
mask_emb_weight.requires_grad = True | |
self.mask_emb = torch.nn.init.normal_(mask_emb_weight, mean=0, std=0.1) | |
self.input_layer_norm = torch.nn.LayerNorm(input_dim, layer_norm_epsilon) | |
self.encoder = encoder | |
self.encoder_top_n_out = torch.nn.parameter.Parameter( | |
torch.Tensor(num_codebooks, self.encoder.output_size(), num_embeddings) | |
) | |
def forward( | |
self, | |
xs: torch.Tensor, | |
xs_lens: torch.Tensor, | |
text: Optional[torch.Tensor] = None, | |
text_length: Optional[torch.Tensor] = None, | |
): | |
# should support nonstreamming and streamming | |
# TODO(Mddct): streamming future | |
# eg: full attenton and chunk or dynamic chunk training | |
# 1 forward subsampling | |
xs, pos_emb, masks = self._forward_subsampling(xs, xs_lens) | |
unmasked_xs = xs | |
# 2 mask features | |
# 2.0 apply mask | |
masked_xs, masked_masks = self._apply_mask(xs) | |
# 2.1 get nearest embedding | |
target_ids = self._nearest_embedding_idx(unmasked_xs) | |
# 3 forward xxx-formaer block | |
out, out_mask = self._forward_encoder_blocks(masked_xs, masks, pos_emb, masks) | |
# 4 get logits | |
out = out.unsqueeze(1) # [B, 1, T', dim] | |
top_n_out = self.encoder_top_n_out.unsqueeze( | |
0 | |
) # [num_codebooks, dim, num_embeddings] | |
out = torch.matmul(out, top_n_out) # [B, num_codebooks, T', num_embeddings] | |
# 5 compute loss | |
loss = self._compute_loss(out, target_ids, out_mask.squeeze(1) * masked_masks) | |
return {"loss": loss} | |
def _compute_loss( | |
self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | |
): | |
input = input.transpose(1, 3) # [B, num_embeddings, T' num_codebooks] | |
entropy = torch.nn.functional.cross_entropy( | |
input, target, reduction="none" | |
) # [B, T', num_codebooks] | |
# stop gradient for non mask area | |
loss = entropy * mask.unsqueeze(2) | |
return loss.sum() / (mask.sum() * loss.size(2)) | |
def _forward_encoder_blocks( | |
self, | |
xs: torch.Tensor, | |
xs_masks: torch.Tensor, | |
pos_emb: torch.Tensor, | |
mask_pad: torch.Tensor, | |
): | |
masks = xs_masks | |
for layer in self.encoder.encoders: | |
xs, masks, _, _ = layer(xs, xs_masks, pos_emb, mask_pad) | |
if self.encoder.normalize_before: | |
xs = self.encoder.after_norm(xs) | |
# Here we assume the mask is not changed in encoder layers, so just | |
# return the masks before encoder layers, and the masks will be used | |
# for cross attention with decoder later | |
return xs, masks | |
def _nearest_embedding_idx(self, xs: torch.Tensor) -> torch.Tensor: | |
xs = self.input_layer_norm(xs) | |
xs = self.input_dropout(xs) | |
xs = torch.matmul(xs, self.projection.to(xs.device)) | |
B, T, C = xs.size() | |
flattened_input = xs.view(-1, C) | |
embeddings = self.embeddings.to( | |
xs.device | |
) # [num_codebooks, embedding_dim, num_embeddings] | |
# [num_codebooks, B*T, num_embeddings] | |
distance = ( | |
torch.sum(flattened_input**2, dim=1, keepdim=True).unsqueeze(0) | |
+ torch.sum(embeddings**2, dim=1, keepdim=True) | |
- 2 * torch.matmul(flattened_input.unsqueeze(0), embeddings) | |
) | |
out = torch.argmin(distance, dim=-1) # [num_codebooks, B*T] | |
out = out.transpose(0, 1) # [B*T, num_codebooks] | |
return out.reshape(B, T, -1) # [B, T, num_codebooks] | |
def _apply_mask(self, xs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
masks = compute_mask_indices( | |
xs.size()[:-1], | |
self.mask_prob, | |
self.mask_length, | |
self.min_masks, | |
device=xs.device, | |
) | |
masks_expand = masks.unsqueeze(-1) # [B, T, 1] | |
mask_emb = self.mask_emb.to(xs.device).view(1, 1, -1) | |
xs = torch.where(masks_expand, mask_emb, xs) | |
return xs, masks | |
def _forward_subsampling( | |
self, xs: torch.Tensor, xs_lens: torch.Tensor | |
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
T = xs.size(1) | |
masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) | |
if self.encoder.global_cmvn is not None: | |
xs = self.encoder.global_cmvn(xs) | |
xs, pos_emb, masks = self.encoder.embed(xs, masks) | |
return xs, pos_emb, masks | |