MIT-SLS/USAD-Small
011
1# Reference: https://github.com/sooftware/conformer2 3import contextlib4import math5from collections import defaultdict6from typing import Dict, List, Optional, Tuple7 8import torch9import torch.nn.functional as F10from torch import nn11from torch.nn.attention import SDPBackend, sdpa_kernel12 13 14def lengths_to_padding_mask(15 lengths: torch.Tensor, max_len: Optional[int] = None16) -> torch.Tensor:17 """Create padding mask from lengths.18 19 Args:20 lengths: A 1-D tensor of shape (B,).21 max_len: An integer. It will be automatically set to the max value of lengths22 if not given.23 24 Returns:25 A bool tensor of shape (B, max_len), where padded positions are indicated by True.26 """27 batch_size = lengths.size(0)28 max_len = lengths.max().item() if max_len is None else max_len29 seq_range = torch.arange(30 0, max_len, dtype=torch.long, device=lengths.device31 )32 seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)33 lengths_expand = lengths.unsqueeze(1).expand_as(seq_range_expand)34 padding_mask = seq_range_expand >= lengths_expand35 return padding_mask36 37 38class SamePad(nn.Module):39 def __init__(self, kernel_size, causal=False):40 super().__init__()41 if causal:42 self.remove = kernel_size - 143 else:44 self.remove = 1 if kernel_size % 2 == 0 else 045 46 def forward(self, x):47 if self.remove > 0:48 x = x[:, :, : -self.remove]49 return x50 51 52class TransposeLast(nn.Module):53 def __init__(self, deconstruct_idx=None, tranpose_dim=-2):54 super().__init__()55 self.deconstruct_idx = deconstruct_idx56 self.tranpose_dim = tranpose_dim57 58 def forward(self, x):59 if self.deconstruct_idx is not None:60 x = x[self.deconstruct_idx]61 return x.transpose(self.tranpose_dim, -1)62 63 64class Swish(nn.Module):65 def __init__(self):66 super(Swish, self).__init__()67 68 def forward(self, inputs: torch.Tensor) -> torch.Tensor:69 return inputs * inputs.sigmoid()70 71 72class GLU(nn.Module):73 def __init__(self, dim: int) -> None:74 super(GLU, self).__init__()75 self.dim = dim76 77 def forward(self, inputs: torch.Tensor) -> torch.Tensor:78 outputs, gate = inputs.chunk(2, dim=self.dim)79 return outputs * gate.sigmoid()80 81 82class RMSNorm(torch.nn.Module):83 def __init__(self, dim: int, eps: float = 1e-5):84 super().__init__()85 self.eps = eps86 self.weight = nn.Parameter(torch.ones(dim))87 88 def _norm(self, x):89 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)90 91 def forward(self, x):92 output = self._norm(x.float()).type_as(x)93 return output * self.weight94 95 96class ResidualConnectionModule(nn.Module):97 def __init__(98 self,99 module: nn.Module,100 module_factor: float = 1.0,101 input_factor: float = 1.0,102 ):103 super(ResidualConnectionModule, self).__init__()104 self.module = module105 self.module_factor = module_factor106 self.input_factor = input_factor107 108 def forward(self, inputs: torch.Tensor) -> torch.Tensor:109 return (self.module(inputs) * self.module_factor) + (110 inputs * self.input_factor111 )112 113 114class Linear(nn.Module):115 def __init__(116 self, in_features: int, out_features: int, bias: bool = True117 ) -> None:118 super(Linear, self).__init__()119 self.linear = nn.Linear(in_features, out_features, bias=bias)120 nn.init.xavier_uniform_(self.linear.weight)121 if bias:122 nn.init.zeros_(self.linear.bias)123 124 def forward(self, x: torch.Tensor) -> torch.Tensor:125 return self.linear(x)126 127 128class View(nn.Module):129 def __init__(self, shape: tuple, contiguous: bool = False):130 super(View, self).__init__()131 self.shape = shape132 self.contiguous = contiguous133 134 def forward(self, x: torch.Tensor) -> torch.Tensor:135 if self.contiguous:136 x = x.contiguous()137 138 return x.view(*self.shape)139 140 141class Transpose(nn.Module):142 def __init__(self, shape: tuple):143 super(Transpose, self).__init__()144 self.shape = shape145 146 def forward(self, x: torch.Tensor) -> torch.Tensor:147 return x.transpose(*self.shape)148 149 150class FeedForwardModule(nn.Module):151 def __init__(152 self,153 encoder_dim: int = 512,154 expansion_factor: int = 4,155 dropout_p: float = 0.1,156 rms_norm: bool = False,157 ) -> None:158 super(FeedForwardModule, self).__init__()159 self.sequential = nn.Sequential(160 (161 nn.LayerNorm(encoder_dim)162 if not rms_norm163 else RMSNorm(encoder_dim)164 ),165 Linear(encoder_dim, encoder_dim * expansion_factor, bias=True),166 Swish(),167 nn.Dropout(p=dropout_p),168 Linear(encoder_dim * expansion_factor, encoder_dim, bias=True),169 nn.Dropout(p=dropout_p),170 )171 172 def forward(self, inputs: torch.Tensor) -> torch.Tensor:173 return self.sequential(inputs)174 175 176class DepthwiseConv1d(nn.Module):177 def __init__(178 self,179 in_channels: int,180 out_channels: int,181 kernel_size: int,182 stride: int = 1,183 padding: int = 0,184 bias: bool = False,185 ) -> None:186 super(DepthwiseConv1d, self).__init__()187 assert (188 out_channels % in_channels == 0189 ), "out_channels should be constant multiple of in_channels"190 self.conv = nn.Conv1d(191 in_channels=in_channels,192 out_channels=out_channels,193 kernel_size=kernel_size,194 groups=in_channels,195 stride=stride,196 padding=padding,197 bias=bias,198 )199 200 def forward(self, inputs: torch.Tensor) -> torch.Tensor:201 return self.conv(inputs)202 203 204class PointwiseConv1d(nn.Module):205 def __init__(206 self,207 in_channels: int,208 out_channels: int,209 stride: int = 1,210 padding: int = 0,211 bias: bool = True,212 ) -> None:213 super(PointwiseConv1d, self).__init__()214 self.conv = nn.Conv1d(215 in_channels=in_channels,216 out_channels=out_channels,217 kernel_size=1,218 stride=stride,219 padding=padding,220 bias=bias,221 )222 223 def forward(self, inputs: torch.Tensor) -> torch.Tensor:224 return self.conv(inputs)225 226 227class ConformerConvModule(nn.Module):228 def __init__(229 self,230 in_channels: int,231 kernel_size: int = 31,232 expansion_factor: int = 2,233 dropout_p: float = 0.1,234 rms_norm: bool = False,235 ) -> None:236 super(ConformerConvModule, self).__init__()237 assert (238 kernel_size - 1239 ) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding"240 assert (241 expansion_factor == 2242 ), "Currently, Only Supports expansion_factor 2"243 244 self.sequential = nn.Sequential(245 (246 nn.LayerNorm(in_channels)247 if not rms_norm248 else RMSNorm(in_channels)249 ),250 Transpose(shape=(1, 2)),251 PointwiseConv1d(252 in_channels,253 in_channels * expansion_factor,254 stride=1,255 padding=0,256 bias=True,257 ),258 GLU(dim=1),259 DepthwiseConv1d(260 in_channels,261 in_channels,262 kernel_size,263 stride=1,264 padding=(kernel_size - 1) // 2,265 ),266 nn.BatchNorm1d(in_channels),267 Swish(),268 PointwiseConv1d(269 in_channels, in_channels, stride=1, padding=0, bias=True270 ),271 nn.Dropout(p=dropout_p),272 )273 274 def forward(self, inputs: torch.Tensor) -> torch.Tensor:275 return self.sequential(inputs).transpose(1, 2)276 277 278class FramewiseConv2dSubampling(nn.Module):279 def __init__(self, out_channels: int, subsample_rate: int = 2) -> None:280 super(FramewiseConv2dSubampling, self).__init__()281 assert subsample_rate in {2, 4}, "subsample_rate should be 2 or 4"282 self.subsample_rate = subsample_rate283 self.cnn = nn.Sequential(284 nn.Conv2d(1, out_channels, kernel_size=3, stride=2),285 nn.ReLU(),286 nn.Conv2d(287 out_channels,288 out_channels,289 kernel_size=3,290 stride=(2 if subsample_rate == 4 else 1, 2),291 padding=(0 if subsample_rate == 4 else 1, 0),292 ),293 nn.ReLU(),294 )295 296 def forward(297 self, inputs: torch.Tensor, input_lengths: torch.Tensor298 ) -> Tuple[torch.Tensor, torch.Tensor]:299 # inputs: (B, T, C) -> (B, 1, T, C)300 if self.subsample_rate == 2 and inputs.shape[1] % 2 == 0:301 inputs = F.pad(inputs, (0, 0, 0, 1), "constant", 0)302 if self.subsample_rate == 4 and inputs.shape[1] % 4 < 3:303 inputs = F.pad(304 inputs, (0, 0, 0, 3 - inputs.shape[1] % 4), "constant", 0305 )306 outputs = self.cnn(inputs.unsqueeze(1))307 batch_size, channels, subsampled_lengths, sumsampled_dim = (308 outputs.size()309 )310 311 outputs = outputs.permute(0, 2, 1, 3)312 outputs = outputs.contiguous().view(313 batch_size, subsampled_lengths, channels * sumsampled_dim314 )315 316 if self.subsample_rate == 4:317 output_lengths = input_lengths >> 2318 else:319 output_lengths = input_lengths >> 1320 321 return outputs, output_lengths322 323 def get_out_dim(self, input_dim: int) -> int:324 # dummy input to get the output dimension325 with torch.no_grad():326 device = next(self.parameters()).device327 inputs = torch.zeros(1, 16, input_dim, device=device)328 input_lengths = torch.tensor([16], device=device)329 outputs, _ = self.forward(inputs, input_lengths)330 return outputs.size(-1)331 332 333class PatchwiseConv2dSubampling(nn.Module):334 def __init__(335 self,336 mel_dim: int,337 out_channels: int,338 patch_size_time: int = 16,339 patch_size_freq: int = 16,340 ) -> None:341 super(PatchwiseConv2dSubampling, self).__init__()342 343 self.mel_dim = mel_dim344 self.patch_size_time = patch_size_time345 self.patch_size_freq = patch_size_freq346 347 self.proj = nn.Conv2d(348 1,349 out_channels,350 kernel_size=(patch_size_time, patch_size_freq),351 stride=(patch_size_time, patch_size_freq),352 padding=0,353 )354 self.cnn = nn.Sequential(355 nn.Conv2d(356 out_channels, out_channels, kernel_size=3, stride=1, padding=1357 ),358 nn.ReLU(),359 nn.Conv2d(360 out_channels, out_channels, kernel_size=3, stride=1, padding=1361 ),362 nn.ReLU(),363 )364 365 @property366 def subsample_rate(self) -> int:367 return self.patch_size_time * self.patch_size_freq // self.mel_dim368 369 def forward(370 self, inputs: torch.Tensor, input_lengths: torch.Tensor371 ) -> Tuple[torch.Tensor, torch.Tensor]:372 assert (373 inputs.shape[2] == self.mel_dim374 ), "inputs.shape[2] should be equal to mel_dim"375 376 # inputs: (B, Time, Freq) -> (B, 1, Time, Freq)377 outputs = self.proj(inputs.unsqueeze(1))378 outputs = self.cnn(outputs)379 # (B, channels, Time // patch_size_time, Freq // patch_size_freq)380 outputs = outputs.flatten(2, 3).transpose(1, 2)381 # (B, (Time // patch_size_time) * (Freq // patch_size_freq), channels)382 383 output_lengths = (384 input_lengths385 // self.patch_size_time386 * (self.mel_dim // self.patch_size_freq)387 )388 389 return outputs, output_lengths390 391 392class RelPositionalEncoding(nn.Module):393 def __init__(self, d_model: int) -> None:394 super(RelPositionalEncoding, self).__init__()395 self.d_model = d_model396 self.pe = None397 398 def extend_pe(self, x: torch.Tensor) -> None:399 if self.pe is not None:400 if self.pe.size(1) >= x.size(1) * 2 - 1:401 if self.pe.dtype != x.dtype or self.pe.device != x.device:402 self.pe = self.pe.to(dtype=x.dtype, device=x.device)403 return404 405 length = x.size(1)406 pe_positive = torch.zeros(length, self.d_model, device="cpu")407 pe_negative = torch.zeros(length, self.d_model, device="cpu")408 position = torch.arange(409 0, length, dtype=torch.float32, device="cpu"410 ).unsqueeze(1)411 div_term = torch.exp(412 torch.arange(0, self.d_model, 2, dtype=torch.float32, device="cpu")413 * -(math.log(10000.0) / self.d_model)414 )415 pe_positive[:, 0::2] = torch.sin(position * div_term)416 pe_positive[:, 1::2] = torch.cos(position * div_term)417 pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)418 pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)419 420 pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)421 pe_negative = pe_negative[1:].unsqueeze(0)422 pe = torch.cat([pe_positive, pe_negative], dim=1)423 self.pe = pe.to(device=x.device, dtype=x.dtype)424 425 def forward(self, x: torch.Tensor) -> torch.Tensor:426 # x: (B, T, C)427 self.extend_pe(x)428 assert self.pe is not None429 pos_emb = self.pe[430 :,431 self.pe.size(1) // 2432 - x.size(1)433 + 1 : self.pe.size(1) // 2434 + x.size(1),435 ]436 return pos_emb437 438 439class RelativeMultiHeadAttention(nn.Module):440 def __init__(441 self,442 d_model: int = 512,443 num_heads: int = 16,444 dropout_p: float = 0.1,445 ):446 super(RelativeMultiHeadAttention, self).__init__()447 assert d_model % num_heads == 0, "d_model % num_heads should be zero."448 self.d_model = d_model449 self.d_head = int(d_model / num_heads)450 self.num_heads = num_heads451 self.sqrt_dim = math.sqrt(self.d_head)452 453 self.query_proj = Linear(d_model, d_model)454 self.key_proj = Linear(d_model, d_model)455 self.value_proj = Linear(d_model, d_model)456 self.pos_proj = Linear(d_model, d_model, bias=False)457 458 self.dropout = nn.Dropout(p=dropout_p)459 self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))460 self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))461 torch.nn.init.xavier_uniform_(self.u_bias)462 torch.nn.init.xavier_uniform_(self.v_bias)463 464 self.out_proj = Linear(d_model, d_model)465 466 @staticmethod467 def _relative_shift(pos_score: torch.Tensor) -> torch.Tensor:468 # pos_score: (B, H, T, 2T-1)469 B, H, T, L = pos_score.size()470 471 # Pad on the left of the last dimension: (B, H, T, 2T)472 pos_score = F.pad(pos_score, (1, 0))473 474 # Reshape to (B, H, 2T, T)475 pos_score = pos_score.view(B, H, L + 1, T)476 477 # Slice and reshape back to (B, H, T, 2T-1)478 pos_score = pos_score[:, :, 1:].view(B, H, T, L)479 480 # Keep only first T positions => (B, H, T, T)481 return pos_score[:, :, :, : (L // 2 + 1)]482 483 def forward(484 self,485 query: torch.Tensor,486 key: torch.Tensor,487 value: torch.Tensor,488 pos_embedding: torch.Tensor,489 padding_mask: Optional[torch.Tensor] = None,490 *,491 need_weights: bool = False,492 use_sdpa: Optional[bool] = None,493 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:494 """495 - If need_weights=True: returns (output, attn) like your original code.496 - If need_weights=False: returns (output, None) and uses SDPA in eval for speed/memory.497 """498 B, Tq, _ = query.size()499 _, Tk, _ = key.size()500 501 # Project502 q = self.query_proj(query) # (B, Tq, C)503 k = self.key_proj(key) # (B, Tk, C)504 v = self.value_proj(value) # (B, Tk, C)505 506 # Reshape to (B, H, T, Dh)507 q = q.view(B, Tq, self.num_heads, self.d_head).transpose(508 1, 2509 ) # (B,H,Tq,Dh)510 k = k.view(B, Tk, self.num_heads, self.d_head).transpose(511 1, 2512 ) # (B,H,Tk,Dh)513 v = v.view(B, Tk, self.num_heads, self.d_head).transpose(514 1, 2515 ) # (B,H,Tk,Dh)516 517 # Positional projection.518 # IMPORTANT: allow pos_embedding to be (1, 2T-1, C) and broadcast across batch.519 # pos_embedding expected length: 2Tq - 1 for self-attn.520 pB = pos_embedding.size(0)521 p = self.pos_proj(pos_embedding) # (pB, L, C)522 p = p.view(pB, -1, self.num_heads, self.d_head).transpose(523 1, 2524 ) # (pB,H,L,Dh)525 526 # Compute position-based bias (scaled) to feed SDPA or add to scores527 # q_pos: (B,H,Tq,Dh), p^T: (pB,H,Dh,L) -> broadcast on pB if pB==1528 q_pos = q + self.v_bias.unsqueeze(0).unsqueeze(2) # (B,H,Tq,Dh)529 pos_score = torch.matmul(q_pos, p.transpose(-2, -1)) # (B,H,Tq,L)530 pos_bias = self._relative_shift(pos_score) # (B,H,Tq,Tq) for self-attn531 pos_bias = pos_bias.mul(532 1.0 / self.sqrt_dim533 ) # scale matches SDPA scaling534 535 if padding_mask is not None:536 # padding_mask: (B, T) -> (B, 1, 1, T) to broadcast with pos_bias (B, H, Tq, Tk)537 # This masks out key positions that are padded across all heads and queries538 if padding_mask.dtype != torch.bool:539 padding_mask = padding_mask.to(torch.bool)540 pos_bias = pos_bias.masked_fill(541 padding_mask[:, None, None, :], -1e9542 )543 544 if use_sdpa is None:545 use_sdpa = (not self.training) and (not need_weights)546 547 # ---- Fast inference path: no attention matrix materialized ----548 if use_sdpa:549 # Content term uses u_bias550 q_content = q + self.u_bias.unsqueeze(0).unsqueeze(551 2552 ) # (B,H,Tq,Dh)553 554 with sdpa_kernel(555 [556 SDPBackend.FLASH_ATTENTION,557 SDPBackend.EFFICIENT_ATTENTION,558 SDPBackend.MATH,559 ]560 ):561 out = F.scaled_dot_product_attention(562 q_content, # (B,H,Tq,Dh)563 k, # (B,H,Tk,Dh)564 v, # (B,H,Tk,Dh)565 attn_mask=pos_bias, # (B,H,Tq,Tk) additive bias566 dropout_p=0.0, # dropout disabled in inference567 is_causal=False,568 ) # (BH, Tq, Dh)569 570 out = out.transpose(1, 2).contiguous().view(B, Tq, self.d_model)571 572 return self.out_proj(out), None573 574 # ---- Reference path (training / if you need attn weights): matches your math ----575 q_content = q + self.u_bias.unsqueeze(0).unsqueeze(2) # (B,H,Tq,Dh)576 content_score = torch.matmul(577 q_content, k.transpose(-2, -1)578 ) # (B,H,Tq,Tk)579 content_score = content_score.mul(1.0 / self.sqrt_dim)580 581 score = content_score + pos_bias # already scaled582 583 attn = F.softmax(score, dim=-1)584 attn = self.dropout(attn)585 586 context = torch.matmul(attn, v) # (B,H,Tq,Dh)587 context = (588 context.transpose(1, 2).contiguous().view(B, Tq, self.d_model)589 )590 591 return self.out_proj(context), attn592 593 594class MultiHeadedSelfAttentionModule(nn.Module):595 def __init__(596 self,597 d_model: int,598 num_heads: int,599 dropout_p: float = 0.1,600 rms_norm: bool = False,601 ):602 super(MultiHeadedSelfAttentionModule, self).__init__()603 self.positional_encoding = RelPositionalEncoding(d_model)604 self.layer_norm = (605 nn.LayerNorm(d_model) if not rms_norm else RMSNorm(d_model)606 )607 self.attention = RelativeMultiHeadAttention(608 d_model, num_heads, dropout_p609 )610 self.dropout = nn.Dropout(p=dropout_p)611 612 def forward(613 self,614 inputs: torch.Tensor,615 padding_mask: Optional[torch.Tensor] = None,616 pos_embedding: Optional[torch.Tensor] = None,617 ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:618 if pos_embedding is None:619 pos_embedding = self.positional_encoding(inputs)620 621 inputs = self.layer_norm(inputs)622 outputs, attn = self.attention(623 inputs,624 inputs,625 inputs,626 pos_embedding=pos_embedding,627 padding_mask=padding_mask,628 )629 630 return self.dropout(outputs), attn, pos_embedding631 632 633class ConformerBlock(nn.Module):634 def __init__(635 self,636 encoder_dim: int = 512,637 attention_type: str = "mhsa",638 num_attention_heads: int = 8,639 feed_forward_expansion_factor: int = 4,640 conv_expansion_factor: int = 2,641 feed_forward_dropout_p: float = 0.1,642 attention_dropout_p: float = 0.1,643 conv_dropout_p: float = 0.1,644 conv_kernel_size: int = 31,645 half_step_residual: bool = True,646 transformer_style: bool = False,647 usad_v2: bool = False,648 pre_norm: bool = False,649 rms_norm: bool = False,650 ):651 super(ConformerBlock, self).__init__()652 653 self.transformer_style = transformer_style654 self.attention_type = attention_type655 self.usad_v2 = usad_v2656 self.pre_norm = pre_norm657 658 if half_step_residual and not transformer_style:659 self.feed_forward_residual_factor = 0.5660 else:661 self.feed_forward_residual_factor = 1662 663 assert (664 attention_type == "mhsa"665 ), "Only 'mhsa' attention is supported in this implementation."666 attention = MultiHeadedSelfAttentionModule(667 d_model=encoder_dim,668 num_heads=num_attention_heads,669 dropout_p=attention_dropout_p,670 rms_norm=rms_norm,671 )672 673 self.ffn_1 = FeedForwardModule(674 encoder_dim=encoder_dim,675 expansion_factor=feed_forward_expansion_factor,676 dropout_p=feed_forward_dropout_p,677 rms_norm=rms_norm,678 )679 self.attention = attention680 if not transformer_style:681 self.conv = ConformerConvModule(682 in_channels=encoder_dim,683 kernel_size=conv_kernel_size,684 expansion_factor=conv_expansion_factor,685 dropout_p=conv_dropout_p,686 rms_norm=rms_norm,687 )688 self.ffn_2 = FeedForwardModule(689 encoder_dim=encoder_dim,690 expansion_factor=feed_forward_expansion_factor,691 dropout_p=feed_forward_dropout_p,692 rms_norm=rms_norm,693 )694 self.layernorm = (695 (696 nn.LayerNorm(encoder_dim)697 if not rms_norm698 else RMSNorm(encoder_dim)699 )700 if not pre_norm701 else nn.Identity()702 )703 704 def forward_attention(705 self,706 x: torch.Tensor,707 pos_embedding: Optional[torch.Tensor] = None,708 padding_mask: Optional[torch.Tensor] = None,709 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:710 attn_out, attn, pos_embedding = self.attention(711 x, pos_embedding=pos_embedding, padding_mask=padding_mask712 )713 return attn_out, attn, pos_embedding714 715 def forward_legacy(716 self,717 x: torch.Tensor,718 pos_embedding: Optional[torch.Tensor] = None,719 padding_mask: Optional[torch.Tensor] = None,720 ) -> Tuple[torch.Tensor, Dict[str, Optional[torch.Tensor]]]:721 # FFN 1722 ffn_1_out = self.ffn_1(x)723 x = ffn_1_out * self.feed_forward_residual_factor + x724 725 # Attention726 attn_out, attn, pos_embedding = self.forward_attention(727 x, pos_embedding, padding_mask728 )729 x = attn_out + x730 731 if self.transformer_style:732 x = self.layernorm(x)733 return x, {734 "ffn_1": ffn_1_out,735 "attn": attn,736 "conv": None,737 "ffn_2": None,738 }739 740 # Convolution741 conv_out = self.conv(x)742 x = conv_out + x743 744 # FFN 2745 ffn_2_out = self.ffn_2(x)746 x = ffn_2_out * self.feed_forward_residual_factor + x747 x = self.layernorm(x)748 749 other = {750 "ffn_1": ffn_1_out,751 "attn": attn,752 "conv": conv_out,753 "ffn_2": ffn_2_out,754 "pos_embedding": pos_embedding,755 }756 757 return x, other758 759 def forward_transformer(760 self,761 x: torch.Tensor,762 pos_embedding: Optional[torch.Tensor] = None,763 padding_mask: Optional[torch.Tensor] = None,764 ) -> Tuple[torch.Tensor, Dict[str, Optional[torch.Tensor]]]:765 # Attention766 attn_out, attn, pos_embedding = self.forward_attention(767 x, pos_embedding, padding_mask768 )769 x = attn_out + x770 771 # FFN772 ffn_out = self.ffn_1(x)773 x = ffn_out * self.feed_forward_residual_factor + x774 775 x = self.layernorm(x)776 return x, {777 "ffn_1": ffn_out,778 "attn": attn,779 "conv": None,780 "ffn_2": None,781 "pos_embedding": pos_embedding,782 }783 784 def forward_conformer(785 self,786 x: torch.Tensor,787 pos_embedding: Optional[torch.Tensor] = None,788 padding_mask: Optional[torch.Tensor] = None,789 ) -> Tuple[torch.Tensor, Dict[str, Optional[torch.Tensor]]]:790 # FFN 1791 ffn_1_out = self.ffn_1(x)792 x = ffn_1_out * self.feed_forward_residual_factor + x793 794 # Attention795 attn_out, attn, pos_embedding = self.forward_attention(796 x, pos_embedding, padding_mask797 )798 x = attn_out + x799 800 # Convolution801 conv_out = self.conv(x)802 x = conv_out + x803 804 # FFN 2805 ffn_2_out = self.ffn_2(x)806 x = ffn_2_out * self.feed_forward_residual_factor + x807 x = self.layernorm(x)808 809 other = {810 "ffn_1": ffn_1_out,811 "attn": attn,812 "conv": conv_out,813 "ffn_2": ffn_2_out,814 "pos_embedding": pos_embedding,815 }816 817 return x, other818 819 def forward(820 self,821 x: torch.Tensor,822 pos_embedding: Optional[torch.Tensor] = None,823 padding_mask: Optional[torch.Tensor] = None,824 ) -> Tuple[torch.Tensor, Dict[str, Optional[torch.Tensor]]]:825 if not self.usad_v2:826 return self.forward_legacy(x, pos_embedding, padding_mask)827 828 if self.transformer_style:829 return self.forward_transformer(x, pos_embedding, padding_mask)830 831 return self.forward_conformer(x, pos_embedding, padding_mask)832 833 834class ConformerEncoder(nn.Module):835 def __init__(self, cfg):836 super(ConformerEncoder, self).__init__()837 838 self.cfg = cfg839 self.framewise_subsample = None840 self.patchwise_subsample = None841 self.framewise_in_proj = None842 self.patchwise_in_proj = None843 assert (844 cfg.use_framewise_subsample or cfg.use_patchwise_subsample845 ), "At least one subsampling method should be used"846 if cfg.use_framewise_subsample:847 self.framewise_subsample = FramewiseConv2dSubampling(848 out_channels=cfg.conv_subsample_channels,849 subsample_rate=cfg.conv_subsample_rate,850 )851 self.framewise_in_proj = nn.Sequential(852 Linear(853 self.framewise_subsample.get_out_dim(cfg.input_dim),854 cfg.encoder_dim,855 ),856 nn.Dropout(p=cfg.input_dropout_p),857 )858 if cfg.use_patchwise_subsample:859 self.patchwise_subsample = PatchwiseConv2dSubampling(860 mel_dim=cfg.input_dim,861 out_channels=cfg.conv_subsample_channels,862 patch_size_time=cfg.patch_size_time,863 patch_size_freq=cfg.patch_size_freq,864 )865 self.patchwise_in_proj = nn.Sequential(866 Linear(867 cfg.conv_subsample_channels,868 cfg.encoder_dim,869 ),870 nn.Dropout(p=cfg.input_dropout_p),871 )872 assert not cfg.use_framewise_subsample or (873 cfg.conv_subsample_rate874 == self.patchwise_subsample.subsample_rate875 ), (876 f"conv_subsample_rate ({cfg.conv_subsample_rate}) != patchwise_subsample.subsample_rate"877 f"({self.patchwise_subsample.subsample_rate})"878 )879 880 self.framewise_norm, self.patchwise_norm = None, None881 if getattr(cfg, "subsample_normalization", False):882 if cfg.use_framewise_subsample:883 self.framewise_norm = (884 nn.LayerNorm(cfg.encoder_dim)885 if not getattr(cfg, "rms_norm", False)886 else RMSNorm(cfg.encoder_dim)887 )888 if cfg.use_patchwise_subsample:889 self.patchwise_norm = (890 nn.LayerNorm(cfg.encoder_dim)891 if not getattr(cfg, "rms_norm", False)892 else RMSNorm(cfg.encoder_dim)893 )894 895 self.conv_pos = None896 self.conv_pos_post_ln = None897 if cfg.conv_pos:898 num_pos_layers = cfg.conv_pos_depth899 k = max(3, cfg.conv_pos_width // num_pos_layers)900 self.conv_pos = nn.Sequential(901 TransposeLast(),902 *[903 nn.Sequential(904 nn.Conv1d(905 cfg.encoder_dim,906 cfg.encoder_dim,907 kernel_size=k,908 padding=k // 2,909 groups=cfg.conv_pos_groups,910 ),911 SamePad(k),912 TransposeLast(),913 nn.LayerNorm(914 cfg.encoder_dim, elementwise_affine=False915 ),916 TransposeLast(),917 nn.GELU(),918 )919 for _ in range(num_pos_layers)920 ],921 TransposeLast(),922 )923 self.conv_pos_post_ln = (924 (925 nn.LayerNorm(cfg.encoder_dim)926 if not getattr(cfg, "rms_norm", False)927 else RMSNorm(cfg.encoder_dim)928 )929 if not getattr(cfg, "pre_norm", False)930 else nn.Identity()931 )932 933 self.layers = nn.ModuleList(934 [935 ConformerBlock(936 encoder_dim=cfg.encoder_dim,937 attention_type=cfg.attention_type,938 num_attention_heads=cfg.num_attention_heads,939 feed_forward_expansion_factor=cfg.feed_forward_expansion_factor,940 conv_expansion_factor=cfg.conv_expansion_factor,941 feed_forward_dropout_p=cfg.feed_forward_dropout_p,942 attention_dropout_p=cfg.attention_dropout_p,943 conv_dropout_p=cfg.conv_dropout_p,944 conv_kernel_size=cfg.conv_kernel_size,945 half_step_residual=cfg.half_step_residual,946 transformer_style=getattr(cfg, "transformer_style", False),947 usad_v2=getattr(cfg, "usad_v2", False),948 pre_norm=getattr(cfg, "pre_norm", False),949 rms_norm=getattr(cfg, "rms_norm", False),950 )951 for _ in range(cfg.num_layers)952 ]953 )954 self.layerdrop_p = getattr(cfg, "layerdrop_p", 0.0)955 956 if cfg.attention_type == "mhsa" and len(self.layers) > 0:957 # Share positional encoding across layers958 shared_pos = None959 for layer in self.layers:960 if isinstance(layer.attention, MultiHeadedSelfAttentionModule):961 if shared_pos is None:962 shared_pos = layer.attention.positional_encoding963 else:964 layer.attention.positional_encoding = shared_pos965 if shared_pos is not None:966 # precompute positional encodings967 # expecting most mel inputs to be fewer than 2000 frames (20 seconds)968 max_len = 2000 // cfg.conv_subsample_rate969 shared_pos.extend_pe(torch.tensor(0.0).expand(1, max_len))970 971 def count_parameters(self) -> int:972 """Count parameters of encoder"""973 return sum([p.numel() for p in self.parameters() if p.requires_grad])974 975 def update_dropout(self, dropout_p: float) -> None:976 """Update dropout probability of encoder"""977 for name, child in self.named_children():978 if isinstance(child, nn.Dropout):979 child.p = dropout_p980 981 def forward(982 self,983 inputs: torch.Tensor,984 input_lengths: Optional[torch.Tensor] = None,985 padding_mask: Optional[torch.Tensor] = None,986 *,987 return_hidden: bool = False,988 freeze_input_layers: bool = False,989 target_layer: Optional[int] = None,990 ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, List[torch.Tensor]]]:991 if input_lengths is None:992 input_lengths = torch.full(993 (inputs.size(0),),994 inputs.size(1),995 dtype=torch.long,996 device=inputs.device,997 )998 999 with (1000 torch.no_grad() if freeze_input_layers else contextlib.ExitStack()1001 ):1002 frame_feat, patch_feat = None, None1003 frame_lengths, patch_lengths = None, None1004 if self.framewise_subsample is not None:1005 assert self.framewise_in_proj is not None1006 frame_feat, frame_lengths = self.framewise_subsample(1007 inputs, input_lengths1008 )1009 frame_feat = self.framewise_in_proj(frame_feat)1010 if self.framewise_norm is not None:1011 frame_feat = self.framewise_norm(frame_feat)1012 1013 if self.patchwise_subsample is not None:1014 assert self.patchwise_in_proj is not None1015 patch_feat, patch_lengths = self.patchwise_subsample(1016 inputs, input_lengths1017 )1018 patch_feat = self.patchwise_in_proj(patch_feat)1019 if self.patchwise_norm is not None:1020 patch_feat = self.patchwise_norm(patch_feat)1021 1022 assert frame_feat is not None or patch_feat is not None1023 assert frame_lengths is not None or patch_lengths is not None1024 1025 if frame_feat is not None and patch_feat is not None:1026 assert frame_lengths is not None and patch_lengths is not None1027 min_len = min(frame_feat.size(1), patch_feat.size(1))1028 frame_feat = frame_feat[:, :min_len]1029 patch_feat = patch_feat[:, :min_len]1030 1031 features = frame_feat + patch_feat1032 output_lengths = (1033 frame_lengths1034 if frame_lengths.max().item() < patch_lengths.max().item()1035 else patch_lengths1036 )1037 elif frame_feat is not None:1038 features = frame_feat1039 output_lengths = frame_lengths1040 else:1041 features = patch_feat1042 output_lengths = patch_lengths1043 1044 assert features is not None1045 assert output_lengths is not None1046 1047 # Positional encoding with convolutional layers1048 if self.conv_pos is not None and self.conv_pos_post_ln is not None:1049 pos = self.conv_pos(features)1050 if not self.training:1051 features = features.add_(pos)1052 else:1053 features = features + pos1054 features = self.conv_pos_post_ln(features)1055 1056 # Create padding mask for attention1057 if padding_mask is not None:1058 # downsample to match features length1059 input_len = padding_mask.size(1)1060 feat_len = features.size(1)1061 factor = input_len / feat_len1062 indices = (1063 torch.arange(feat_len, device=padding_mask.device) * factor1064 ).long()1065 padding_mask = padding_mask.index_select(1, indices)1066 else:1067 # create from output_lengths1068 padding_mask = lengths_to_padding_mask(1069 output_lengths, max_len=features.size(1)1070 )1071 1072 layer_results = defaultdict(list)1073 outputs = features1074 other = {}1075 for i, layer in enumerate(self.layers):1076 if (1077 self.training1078 and self.layerdrop_p > 01079 and torch.rand(1).item() < self.layerdrop_p1080 ):1081 continue1082 outputs, other = layer(1083 outputs,1084 pos_embedding=other.get("pos_embedding"),1085 padding_mask=padding_mask,1086 )1087 if return_hidden:1088 layer_results["hidden_states"].append(outputs)1089 for k, v in other.items():1090 layer_results[k].append(v)1091 1092 if target_layer is not None and i + 1 == target_layer:1093 break1094 1095 return outputs, output_lengths, layer_results1096