File size: 7,774 Bytes
ce3c8df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """
Zipformer-inspired encoder (arXiv:2310.11230): a U-Net-shaped stack that
downsamples the middle layers in time and upsamples back before the output
heads, with Conformer-style macaron-FF + self-attention + conv blocks.
Built from standard PyTorch primitives, not a byte-exact port of icefall's
custom ops.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional
from src.scaling import (
ConvolutionModule,
FeedForwardModule,
PositionalEncoding,
SelfAttentionModule,
conv_subsample_length,
)
def make_pad_mask(lengths: torch.Tensor, max_len: int) -> torch.Tensor:
"""Returns (B, max_len) bool tensor, True at padded positions."""
batch_size = lengths.size(0)
seq_range = torch.arange(0, max_len, device=lengths.device).unsqueeze(0).expand(batch_size, -1)
return seq_range >= lengths.unsqueeze(1)
class Conv2dSubsampling(nn.Module):
"""4x time subsampling: two stride-2 conv2d layers turn (T, n_mels) into a
d_model sequence at 1/4 the frame rate (10ms -> 40ms per frame).
"""
def __init__(self, n_mels: int, d_model: int, conv_channels: int = 128, dropout: float = 0.1):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, conv_channels, kernel_size=3, stride=2, padding=1),
nn.SiLU(),
nn.Conv2d(conv_channels, conv_channels, kernel_size=3, stride=2, padding=1),
nn.SiLU(),
)
freq_out = n_mels
for _ in range(2):
freq_out = (freq_out + 2 * 1 - 3) // 2 + 1
self.out_proj = nn.Linear(conv_channels * freq_out, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, feats: torch.Tensor, lengths: torch.Tensor):
# feats: (B, T, n_mels)
x = feats.unsqueeze(1) # (B, 1, T, F)
x = self.conv(x) # (B, C, T', F')
b, c, t, f = x.shape
x = x.permute(0, 2, 1, 3).reshape(b, t, c * f) # (B, T', C*F')
x = self.out_proj(x)
x = self.dropout(x)
out_lengths = conv_subsample_length(lengths, num_layers=2, kernel_size=3, stride=2, padding=1)
out_lengths = out_lengths.clamp(min=0, max=t)
return x, out_lengths
class DownsampleConcat(nn.Module):
"""Groups every `factor` consecutive frames and projects down to d_model.
Unlike a strided conv, no input frame's information is dropped.
"""
def __init__(self, d_model: int, factor: int):
super().__init__()
self.factor = factor
self.proj = nn.Linear(d_model * factor, d_model)
def forward(self, x: torch.Tensor, lengths: torch.Tensor):
b, t, c = x.shape
pad_len = (self.factor - t % self.factor) % self.factor
if pad_len > 0:
x = F.pad(x, (0, 0, 0, pad_len))
t_padded = t + pad_len
x = x.reshape(b, t_padded // self.factor, self.factor * c)
x = self.proj(x)
new_lengths = torch.div(lengths + self.factor - 1, self.factor, rounding_mode="floor")
return x, new_lengths
class UpsampleRepeat(nn.Module):
"""Inverse of DownsampleConcat: repeats each frame `factor` times, projects
back to d_model, then trims/pads to match the skip connection's length.
"""
def __init__(self, d_model: int, factor: int):
super().__init__()
self.factor = factor
self.proj = nn.Linear(d_model, d_model)
def forward(self, x: torch.Tensor, target_time_len: int):
x = x.repeat_interleave(self.factor, dim=1)
x = self.proj(x)
cur_len = x.size(1)
if cur_len > target_time_len:
x = x[:, :target_time_len, :]
elif cur_len < target_time_len:
x = F.pad(x, (0, 0, 0, target_time_len - cur_len))
return x
class ZipformerBlock(nn.Module):
"""One encoder block: macaron FF -> self-attn -> conv -> macaron FF -> norm."""
def __init__(self, d_model: int, nhead: int, d_ff: int, conv_kernel: int, dropout: float):
super().__init__()
self.ff1 = FeedForwardModule(d_model, d_ff, dropout)
self.self_attn = SelfAttentionModule(d_model, nhead, dropout)
self.conv = ConvolutionModule(d_model, conv_kernel, dropout)
self.ff2 = FeedForwardModule(d_model, d_ff, dropout)
self.final_norm = nn.LayerNorm(d_model)
def forward(self, x: torch.Tensor, pad_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
x = self.ff1(x)
x = self.self_attn(x, pad_mask=pad_mask)
x = self.conv(x, pad_mask=pad_mask)
x = self.ff2(x)
x = self.final_norm(x)
return x
class Stage(nn.Module):
def __init__(self, num_layers: int, d_model: int, nhead: int, d_ff: int, conv_kernel: int, dropout: float):
super().__init__()
self.blocks = nn.ModuleList(
[ZipformerBlock(d_model, nhead, d_ff, conv_kernel, dropout) for _ in range(num_layers)]
)
def forward(self, x: torch.Tensor, pad_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
for block in self.blocks:
x = block(x, pad_mask=pad_mask)
return x
class ZipformerEncoder(nn.Module):
"""5-stage U-Net encoder: full -> half -> quarter (bottleneck) -> half ->
full res, with additive skip connections bridging matching resolutions.
"""
def __init__(
self,
n_mels: int = 80,
d_model: int = 256,
conv_channels: int = 128,
nhead: int = 4,
d_ff: int = 1024,
conv_kernel: int = 15,
stage_layers=(2, 3, 4, 3, 2),
downsample_factors=(2, 2),
dropout: float = 0.1,
):
super().__init__()
assert len(stage_layers) == 5, "expected 5 stages: full, half, bottleneck, half, full"
assert len(downsample_factors) == 2
self.frontend = Conv2dSubsampling(n_mels, d_model, conv_channels, dropout)
self.pos_enc = PositionalEncoding(d_model, dropout=dropout)
self.stage0 = Stage(stage_layers[0], d_model, nhead, d_ff, conv_kernel, dropout)
self.down0 = DownsampleConcat(d_model, downsample_factors[0])
self.stage1 = Stage(stage_layers[1], d_model, nhead, d_ff, conv_kernel, dropout)
self.down1 = DownsampleConcat(d_model, downsample_factors[1])
self.stage2 = Stage(stage_layers[2], d_model, nhead, d_ff, conv_kernel, dropout)
self.up1 = UpsampleRepeat(d_model, downsample_factors[1])
self.stage3 = Stage(stage_layers[3], d_model, nhead, d_ff, conv_kernel, dropout)
self.up0 = UpsampleRepeat(d_model, downsample_factors[0])
self.stage4 = Stage(stage_layers[4], d_model, nhead, d_ff, conv_kernel, dropout)
self.final_norm = nn.LayerNorm(d_model)
def forward(self, feats: torch.Tensor, lengths: torch.Tensor):
"""
feats: (B, T, n_mels) log-mel features
lengths: (B,) valid frame counts before padding
returns: (encoder_out (B, T0, d_model), out_lengths (B,))
"""
x, lengths0 = self.frontend(feats, lengths)
x = self.pos_enc(x)
t0 = x.size(1)
pad_mask0 = make_pad_mask(lengths0, t0)
x = self.stage0(x, pad_mask0)
skip0 = x
x, lengths1 = self.down0(x, lengths0)
t1 = x.size(1)
pad_mask1 = make_pad_mask(lengths1, t1)
x = self.stage1(x, pad_mask1)
skip1 = x
x, lengths2 = self.down1(x, lengths1)
t2 = x.size(1)
pad_mask2 = make_pad_mask(lengths2, t2)
x = self.stage2(x, pad_mask2)
x = self.up1(x, target_time_len=t1)
x = x + skip1
x = self.stage3(x, pad_mask1)
x = self.up0(x, target_time_len=t0)
x = x + skip0
x = self.stage4(x, pad_mask0)
x = self.final_norm(x)
return x, lengths0
|