Spaces:
Sleeping
Sleeping
File size: 7,308 Bytes
82e11f0 | 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 | import copy
from typing import List, Optional, Tuple, Union
import torch
from torch import Tensor, nn
from .encoder_layer import BypassModule
from .scaling import FloatLike
class DownsampledZipformer2Encoder(nn.Module):
r"""
DownsampledZipformer2Encoder is a zipformer encoder evaluated at a reduced frame rate,
after convolutional downsampling, and then upsampled again at the output, and combined
with the origin input, so that the output has the same shape as the input.
"""
def __init__(self, encoder: nn.Module, dim: int, downsample: int, dropout: FloatLike) -> None:
super(DownsampledZipformer2Encoder, self).__init__()
self.downsample_factor = downsample
self.downsample = SimpleDownsample(dim, downsample, dropout)
self.num_layers = encoder.num_layers
self.encoder = encoder
self.upsample = SimpleUpsample(dim, downsample)
self.out_combiner = BypassModule(dim, straight_through_rate=0)
def forward(
self,
src: Tensor,
chunk_size: int = -1,
feature_mask: Union[Tensor, float] = 1.0,
attn_mask: Optional[Tensor] = None,
src_key_padding_mask: Optional[Tensor] = None,
) -> Tensor:
r"""Downsample, go through encoder, upsample.
Args:
src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
feature_mask: something that broadcasts with src, that we'll multiply `src`
by at every layer: if a Tensor, likely of shape (seq_len, batch_size, embedding_dim)
attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) or (seq_len, seq_len),
interpreted as (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len).
True means masked position. May be None.
src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); True means
masked position. May be None.
Returns: a Tensor with the same shape as src.
"""
src_orig = src
ds = self.downsample_factor
src = self.downsample(src, src_key_padding_mask)
if attn_mask is not None:
attn_mask = attn_mask[::ds, ::ds]
downsampled_padding_mask = None if src_key_padding_mask is None else src_key_padding_mask[..., ::ds]
src = self.encoder(
src,
chunk_size=chunk_size // ds,
feature_mask=feature_mask,
attn_mask=attn_mask,
src_key_padding_mask=downsampled_padding_mask,
)
src = self.upsample(src)
# remove any extra frames that are not a multiple of downsample_factor
src = src[: src_orig.shape[0]]
return self.out_combiner(src_orig, src)
def streaming_forward(
self,
src: Tensor,
states: List[Tensor],
left_context_len: int,
src_key_padding_mask: Tensor,
) -> Tuple[Tensor, List[Tensor]]:
r"""Downsample, go through encoder, upsample, in streaming forward mode.
Args:
src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
states: list of cached tensors of N encoder layers. For layer-i, states[i*6:(i+1)*6] is
(cached_key, cached_nonlin_attn, cached_val1, cached_val2, cached_conv1, cached_conv2).
left_context_len: Number of left context frames.
src_key_padding_mask: the mask for padding, of shape (batch_size, left_context_len+seq_len);
True means masked position. May be None.
Returns:
- output, a Tensor with the same shape as src.
- updated states
"""
src_orig = src
src = self.downsample(src)
src, new_states = self.encoder.streaming_forward(
src,
states=states,
left_context_len=left_context_len,
src_key_padding_mask=src_key_padding_mask,
)
src = self.upsample(src)
# remove any extra frames that are not a multiple of downsample_factor
src = src[: src_orig.shape[0]]
return self.out_combiner(src_orig, src), new_states
class SimpleDownsample(torch.nn.Module):
"""
Does downsampling with attention, by weighted sum, and a projection..
"""
def __init__(
self,
channels: int,
downsample: int,
dropout: FloatLike,
) -> None:
super(SimpleDownsample, self).__init__()
self.bias = nn.Parameter(torch.zeros(downsample))
self.name = None # will be set from training code
self.dropout = copy.deepcopy(dropout)
self.downsample = downsample
def forward(
self,
src: Tensor,
src_key_padding_mask: Optional[Tensor] = None,
) -> Tensor:
"""
x: (seq_len, batch_size, in_channels)
Returns a tensor of shape
( (seq_len+downsample-1)//downsample, batch_size, channels)
"""
(seq_len, batch_size, in_channels) = src.shape
ds = self.downsample
d_seq_len = (seq_len + ds - 1) // ds
# ---------------------------------------------------------
# Match unpadded inference by extending each sequence's last valid frame.
# ---------------------------------------------------------
if src_key_padding_mask is not None:
lengths = (~src_key_padding_mask).sum(dim=-1)
last_valid_indices = lengths.sub(1).view(1, batch_size, 1)
last_valid = src.gather(
0,
last_valid_indices.expand(1, batch_size, in_channels),
)
src = torch.where(
src_key_padding_mask.transpose(0, 1).unsqueeze(-1),
last_valid,
src,
)
# Pad to an exact multiple of self.downsample
# right-pad src, repeating the last element.
pad = d_seq_len * ds - seq_len
src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2])
src = torch.cat((src, src_extra), dim=0)
assert src.shape[0] == d_seq_len * ds
src = src.reshape(d_seq_len, ds, batch_size, in_channels)
weights = self.bias.softmax(dim=0)
# weights: (downsample, 1, 1)
weights = weights.unsqueeze(-1).unsqueeze(-1)
# ans1 is the first `in_channels` channels of the output
ans = (src * weights).sum(dim=1)
return ans
class SimpleUpsample(torch.nn.Module):
"""
A very simple form of upsampling that mostly just repeats the input, but
also adds a position-specific bias.
"""
def __init__(self, num_channels: int, upsample: int) -> None:
super(SimpleUpsample, self).__init__()
self.upsample = upsample
def forward(self, src: Tensor) -> Tensor:
"""
x: (seq_len, batch_size, num_channels)
Returns a tensor of shape
( (seq_len*upsample), batch_size, num_channels)
"""
upsample = self.upsample
(seq_len, batch_size, num_channels) = src.shape
src = src.unsqueeze(1).expand(seq_len, upsample, batch_size, num_channels)
src = src.reshape(seq_len * upsample, batch_size, num_channels)
return src
|