From 80a14f93d3d7e3de50c3623cc055333128bb751e Mon Sep 17 00:00:00 2001 From: yaozengwei Date: Thu, 20 Jul 2023 19:38:03 +0800 Subject: [PATCH] Use block-wise attention --- egs/librispeech/ASR/zipformer/scaling.py | 27 +++ egs/librispeech/ASR/zipformer/train.py | 8 + egs/librispeech/ASR/zipformer/zipformer.py | 267 ++++++++++++++------- 3 files changed, 215 insertions(+), 87 deletions(-) diff --git a/egs/librispeech/ASR/zipformer/scaling.py b/egs/librispeech/ASR/zipformer/scaling.py index 4ee7b7826..6c315b8e8 100644 --- a/egs/librispeech/ASR/zipformer/scaling.py +++ b/egs/librispeech/ASR/zipformer/scaling.py @@ -1602,6 +1602,33 @@ def convert_num_channels(x: Tensor, num_channels: int) -> Tensor: return torch.cat((x, zeros), dim=-1) +def unfold( + x: Tensor, x_pad: int, num_blocks: int, kernel: int, stride: int, padding: int +) -> Tensor: + """ + Args: + x: input of shape (seq_len, batch_size, channel) + Returns: + blocks: (kernel, batch_size * num_blocks, channel) + """ + seq_len, batch_size, channel = x.size() + x = x.permute(1, 2, 0) # (B, D, T) + + x = nn.functional.pad(x, pad=(0, x_pad), value=0.0) + + blocks = nn.functional.unfold( + x.unsqueeze(-1), + kernel_size=(kernel, 1), + padding=(padding, 0), + stride=(stride, 1), + ) # (B, C * kernel, num_blocks) + blocks = blocks.reshape(batch_size, channel, kernel, num_blocks) + blocks = blocks.permute(2, 0, 3, 1) + blocks = blocks.reshape(kernel, batch_size * num_blocks, channel) + + return blocks + + def _test_whiten(): for proportion in [0.1, 0.5, 10.0]: logging.info(f"_test_whiten(): proportion = {proportion}") diff --git a/egs/librispeech/ASR/zipformer/train.py b/egs/librispeech/ASR/zipformer/train.py index bc3e9c1ba..5a8dae619 100755 --- a/egs/librispeech/ASR/zipformer/train.py +++ b/egs/librispeech/ASR/zipformer/train.py @@ -187,6 +187,13 @@ def add_model_arguments(parser: argparse.ArgumentParser): help="Positional-encoding embedding dimension", ) + parser.add_argument( + "--block-size", + type=int, + default="32", + help="Block size used in block-wise attention", + ) + parser.add_argument( "--encoder-unmasked-dim", type=str, @@ -574,6 +581,7 @@ def get_encoder_model(params: AttributeDict) -> nn.Module: num_heads=_to_int_tuple(params.num_heads), feedforward_dim=_to_int_tuple(params.feedforward_dim), cnn_module_kernel=_to_int_tuple(params.cnn_module_kernel), + block_size=params.block_size, dropout=ScheduledFloat((0.0, 0.3), (20000.0, 0.1)), warmup_batches=4000.0, causal=params.causal, diff --git a/egs/librispeech/ASR/zipformer/zipformer.py b/egs/librispeech/ASR/zipformer/zipformer.py index 7d98dbeb1..d12b9f22b 100644 --- a/egs/librispeech/ASR/zipformer/zipformer.py +++ b/egs/librispeech/ASR/zipformer/zipformer.py @@ -39,6 +39,7 @@ from scaling import ( FloatLike, limit_param_value, convert_num_channels, + unfold, ) from torch import Tensor, nn @@ -105,6 +106,7 @@ class Zipformer2(EncoderInterface): feedforward_dim: Union[int, Tuple[int]] = 1536, cnn_module_kernel: Union[int, Tuple[int]] = 31, pos_dim: int = 192, + block_size: int = 32, dropout: FloatLike = None, # see code below for default warmup_batches: float = 4000.0, causal: bool = False, @@ -140,6 +142,7 @@ class Zipformer2(EncoderInterface): self.num_heads = num_heads = _to_tuple(num_heads) feedforward_dim = _to_tuple(feedforward_dim) self.cnn_module_kernel = cnn_module_kernel = _to_tuple(cnn_module_kernel) + self.block_size = block_size self.causal = causal self.chunk_size = chunk_size @@ -153,6 +156,7 @@ class Zipformer2(EncoderInterface): num_encoders = len(downsampling_factor) for i in range(num_encoders): + ds = downsampling_factor[i] encoder_layer = Zipformer2EncoderLayer( embed_dim=encoder_dim[i], @@ -164,6 +168,7 @@ class Zipformer2(EncoderInterface): feedforward_dim=feedforward_dim[i], dropout=dropout, cnn_module_kernel=cnn_module_kernel[i], + block_size=block_size // ds, causal=causal, ) @@ -173,13 +178,14 @@ class Zipformer2(EncoderInterface): encoder_layer, num_encoder_layers[i], pos_dim=pos_dim, + block_size=block_size // ds, dropout=dropout, warmup_begin=warmup_batches * (i + 1) / (num_encoders + 1), warmup_end=warmup_batches * (i + 2) / (num_encoders + 1), final_layerdrop_rate=0.035 * (downsampling_factor[i] ** 0.5), ) - if downsampling_factor[i] != 1: + if ds != 1: encoder = DownsampledZipformer2Encoder( encoder, dim=encoder_dim[i], @@ -536,6 +542,7 @@ class Zipformer2EncoderLayer(nn.Module): feedforward_dim: int, dropout: FloatLike = 0.1, cnn_module_kernel: int = 31, + block_size: int = 32, causal: bool = False, attention_skip_rate: FloatLike = ScheduledFloat((0.0, 0.2), (4000.0, 0.05), (16000, 0.0), default=0), conv_skip_rate: FloatLike = ScheduledFloat((0.0, 0.2), (4000.0, 0.05), (16000, 0.0), default=0), @@ -569,14 +576,14 @@ class Zipformer2EncoderLayer(nn.Module): self.self_attn_weights = RelPositionMultiheadAttentionWeights( embed_dim, pos_dim=pos_dim, num_heads=num_heads, query_head_dim=query_head_dim, pos_head_dim=pos_head_dim, - dropout=0.0, + block_size=block_size, dropout=0.0, ) self.self_attn1 = SelfAttention(embed_dim, num_heads, - value_head_dim) + value_head_dim, block_size=block_size) self.self_attn2 = SelfAttention(embed_dim, num_heads, - value_head_dim) + value_head_dim, block_size=block_size) self.feed_forward1 = FeedforwardModule(embed_dim, (feedforward_dim * 3) // 4, @@ -591,7 +598,8 @@ class Zipformer2EncoderLayer(nn.Module): dropout) self.nonlin_attention = NonlinAttention(embed_dim, - hidden_channels=3 * embed_dim // 4) + hidden_channels=3 * embed_dim // 4, + block_size=block_size) self.conv_module1 = ConvolutionModule(embed_dim, cnn_module_kernel, @@ -917,6 +925,7 @@ class Zipformer2Encoder(nn.Module): encoder_layer: nn.Module, num_layers: int, pos_dim: int, + block_size: int, dropout: float, warmup_begin: float, warmup_end: float, @@ -931,6 +940,7 @@ class Zipformer2Encoder(nn.Module): [copy.deepcopy(encoder_layer) for i in range(num_layers)] ) self.num_layers = num_layers + self.block_size = block_size assert 0 <= warmup_begin <= warmup_end @@ -966,7 +976,7 @@ class Zipformer2Encoder(nn.Module): Returns: a Tensor with the same shape as src. """ - pos_emb = self.encoder_pos(src) + pos_emb = self.encoder_pos(src, block_size=self.block_size) output = src if not torch.jit.is_scripting() and not torch.jit.is_tracing(): @@ -1314,9 +1324,9 @@ class CompactRelPositionalEncoding(torch.nn.Module): self.length_factor = length_factor self.extend_pe(torch.tensor(0.0).expand(max_len)) - def extend_pe(self, x: Tensor, left_context_len: int = 0) -> None: + def extend_pe(self, x: Tensor) -> None: """Reset the positional encodings.""" - T = x.size(0) + left_context_len + T = x.size(0) if self.pe is not None: # self.pe contains both positive and negative parts @@ -1361,25 +1371,25 @@ class CompactRelPositionalEncoding(torch.nn.Module): self.pe = pe.to(dtype=x.dtype) - def forward(self, x: Tensor, left_context_len: int = 0) -> Tensor: + def forward(self, x: Tensor, block_size: int = 0) -> Tensor: """Create positional encoding. Args: x (Tensor): Input tensor (time, batch, `*`). - left_context_len: (int): Length of cached left context. + block_size (int): Returns: - positional embedding, of shape (batch, left_context_len + 2*time-1, `*`). + positional embedding, of shape (1, 2*time-1, `*`) or (1, 4*block_size-1, `*`). """ - self.extend_pe(x, left_context_len) - x_size_left = x.size(0) + left_context_len - # length of positive side: x.size(0) + left_context_len - # length of negative side: x.size(0) + self.extend_pe(x) + rel_pos = 2 * block_size if block_size != 0 else x.size(0) + # length of positive side: 2 * block_size + # length of negative side: 2 * block_size pos_emb = self.pe[ self.pe.size(0) // 2 - - x_size_left + - rel_pos + 1 : self.pe.size(0) // 2 # noqa E203 - + x.size(0), + + rel_pos, : ] pos_emb = pos_emb.unsqueeze(0) @@ -1413,6 +1423,7 @@ class RelPositionMultiheadAttentionWeights(nn.Module): num_heads: int, query_head_dim: int, pos_head_dim: int, + block_size: int, dropout: float = 0.0, pos_emb_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.0)) @@ -1422,6 +1433,7 @@ class RelPositionMultiheadAttentionWeights(nn.Module): self.num_heads = num_heads self.query_head_dim = query_head_dim self.pos_head_dim = pos_head_dim + self.block_size = block_size self.dropout = dropout self.pos_emb_skip_rate = copy.deepcopy(pos_emb_skip_rate) self.name = None # will be overwritten in training code; for diagnostics. @@ -1478,16 +1490,19 @@ class RelPositionMultiheadAttentionWeights(nn.Module): r""" Args: x: input of shape (seq_len, batch_size, embed_dim) - pos_emb: Positional embedding tensor, of shape (1, 2*seq_len - 1, pos_dim) + pos_emb: Positional embedding tensor, of shape (1, 4*block_size-1, pos_dim) key_padding_mask: a bool tensor of shape (batch_size, seq_len). Positions that are True in this mask will be ignored as sources in the attention weighting. attn_mask: mask of shape (seq_len, seq_len) or (batch_size, seq_len, seq_len), interpreted as ([batch_size,] tgt_seq_len, src_seq_len) saying which positions are allowed to attend to which other positions. Returns: - a tensor of attention weights, of shape (hum_heads, batch_size, seq_len, seq_len) - interpreted as (hum_heads, batch_size, tgt_seq_len, src_seq_len). + a tensor of attention weights, of shape (hum_heads, batch_size * num_blocks, block_size, block_size * 3) + interpreted as (hum_heads, batch_size * num_blocks, tgt_seq_len, src_seq_len), + where num_blocks = (seq_len + block_size - 1) // block_size. """ + assert attn_mask is None, "Not supported yet" + x = self.in_proj(x) query_head_dim = self.query_head_dim pos_head_dim = self.pos_head_dim @@ -1508,16 +1523,31 @@ class RelPositionMultiheadAttentionWeights(nn.Module): k = self.whiten_keys(self.balance_keys(k)) # does nothing in the forward pass. p = self.copy_pos_query(p) # for diagnostics only, does nothing. - q = q.reshape(seq_len, batch_size, num_heads, query_head_dim) - p = p.reshape(seq_len, batch_size, num_heads, pos_head_dim) - k = k.reshape(seq_len, batch_size, num_heads, query_head_dim) + # divide into blocks by unfold function + block_size = self.block_size + num_blocks = (seq_len + block_size - 1) // block_size + pad_len = num_blocks * block_size - seq_len + + # (kernel, batch_size * num_blocks, channel) + q_blocks = unfold(q, pad_len, num_blocks, kernel=block_size, stride=block_size, padding=0) + p_blocks = unfold(p, pad_len, num_blocks, kernel=block_size, stride=block_size, padding=0) + k_blocks = unfold(k, pad_len, num_blocks, kernel=block_size * 3, stride=block_size, padding=block_size) # time1 refers to target, time2 refers to source. - q = q.permute(2, 1, 0, 3) # (head, batch, time1, query_head_dim) - p = p.permute(2, 1, 0, 3) # (head, batch, time1, pos_head_dim) - k = k.permute(2, 1, 3, 0) # (head, batch, d_k, time2) + time1 = q_blocks.size(0) + time2 = k_blocks.size(0) + new_batch_size = batch_size * num_blocks - attn_scores = torch.matmul(q, k) + q_blocks = q_blocks.reshape(time1, new_batch_size, num_heads, query_head_dim) + p_blocks = p_blocks.reshape(time1, new_batch_size, num_heads, pos_head_dim) + k_blocks = k_blocks.reshape(time2, new_batch_size, num_heads, query_head_dim) + + q_blocks = q_blocks.permute(2, 1, 0, 3) # (head, new_batch, time1, query_head_dim) + p_blocks = p_blocks.permute(2, 1, 0, 3) # (head, new_batch, time1, pos_head_dim) + k_blocks = k_blocks.permute(2, 1, 3, 0) # (head, new_batch, d_k, time2) + + # (head, new_batch, time1, time2) + attn_scores = torch.matmul(q_blocks, k_blocks) use_pos_scores = False if torch.jit.is_scripting() or torch.jit.is_tracing(): @@ -1528,32 +1558,21 @@ class RelPositionMultiheadAttentionWeights(nn.Module): if use_pos_scores: pos_emb = self.linear_pos(pos_emb) - seq_len2 = 2 * seq_len - 1 - pos_emb = pos_emb.reshape(-1, seq_len2, num_heads, pos_head_dim).permute(2, 0, 3, 1) - # pos shape now: (head, {1 or batch_size}, pos_dim, seq_len2) + pos_emb = pos_emb.reshape(1, time1 + time2 - 1, num_heads, pos_head_dim).permute(2, 0, 3, 1) + # pos shape now: (head, 1, pos_dim, time1+time2-1) - # (head, batch, time1, pos_dim) x (head, 1, pos_dim, seq_len2) -> (head, batch, time1, seq_len2) + # (head, batch, time1, pos_dim) x (head, 1, pos_dim, time1+time2-1) -> (head, batch, time1, time1+time2-1) # [where seq_len2 represents relative position.] - pos_scores = torch.matmul(p, pos_emb) + pos_scores = torch.matmul(p_blocks, pos_emb) # the following .as_strided() expression converts the last axis of pos_scores from relative # to absolute position. I don't know whether I might have got the time-offsets backwards or # not, but let this code define which way round it is supposed to be. - if torch.jit.is_tracing(): - (num_heads, batch_size, time1, n) = pos_scores.shape - rows = torch.arange(start=time1 - 1, end=-1, step=-1) - cols = torch.arange(seq_len) - rows = rows.repeat(batch_size * num_heads).unsqueeze(-1) - indexes = rows + cols - pos_scores = pos_scores.reshape(-1, n) - pos_scores = torch.gather(pos_scores, dim=1, index=indexes) - pos_scores = pos_scores.reshape(num_heads, batch_size, time1, seq_len) - else: - pos_scores = pos_scores.as_strided((num_heads, batch_size, seq_len, seq_len), - (pos_scores.stride(0), - pos_scores.stride(1), - pos_scores.stride(2)-pos_scores.stride(3), - pos_scores.stride(3)), - storage_offset=pos_scores.stride(3) * (seq_len - 1)) + pos_scores = pos_scores.as_strided((num_heads, new_batch_size, time1, time2), + (pos_scores.stride(0), + pos_scores.stride(1), + pos_scores.stride(2)-pos_scores.stride(3), + pos_scores.stride(3)), + storage_offset=pos_scores.stride(3) * (time1 - 1)) attn_scores = attn_scores + pos_scores @@ -1577,9 +1596,10 @@ class RelPositionMultiheadAttentionWeights(nn.Module): penalty=1.0e-04, name=self.name) - assert attn_scores.shape == (num_heads, batch_size, seq_len, seq_len) + assert attn_scores.shape == (num_heads, new_batch_size, time1, time2) if attn_mask is not None: + # TODO: assert attn_mask.dtype == torch.bool # use -1000 to avoid nan's where attn_mask and key_padding_mask make # all scores zero. It's important that this be large enough that exp(-1000) @@ -1587,12 +1607,31 @@ class RelPositionMultiheadAttentionWeights(nn.Module): # compares the final weights with zero. attn_scores = attn_scores.masked_fill(attn_mask, -1000) - if key_padding_mask is not None: - assert key_padding_mask.shape == (batch_size, seq_len), key_padding_mask.shape - attn_scores = attn_scores.masked_fill( - key_padding_mask.unsqueeze(1), - -1000, - ) + assert key_padding_mask is not None + assert key_padding_mask.shape == (batch_size, seq_len), key_padding_mask.shape + attn_offsets = (~key_padding_mask).float() # 0 at padding positions + + # (seq_len, batch, 1) + attn_offsets = attn_offsets.transpose(0, 1).unsqueeze(-1) + # (time2, new_batch_size) + attn_offsets = unfold( + attn_offsets, pad_len, num_blocks, + kernel=block_size * 3, stride=block_size, padding=block_size, + ).squeeze(-1) + + # For the blocks are all padding + all_pad_mask = (attn_offsets.sum(dim=0, keepdim=True) == 0) # (1, new_batch_size) + all_pad_mask = all_pad_mask.unsqueeze(-1).unsqueeze(-1) # (1, new_batch_size, 1, 1) + + attn_offsets = 1 - attn_offsets # 1 at padding positions + # attn_offsets[attn_offsets != 0] = float("-inf") + attn_offsets[attn_offsets != 0] = -1000 + # attn_offsets = attn_offsets.masked_fill((attn_offsets != 0), -1000) + + # (1, new_batch_size, 1, time2) + attn_offsets = attn_offsets.transpose(0, 1).unsqueeze(1).unsqueeze(0) + + attn_scores = attn_scores + attn_offsets # We use our own version of softmax, defined in scaling.py, which should # save a little of the memory used in backprop by, if we are in @@ -1600,6 +1639,9 @@ class RelPositionMultiheadAttentionWeights(nn.Module): # half-precision output for backprop purposes. attn_weights = softmax(attn_scores, dim=-1) + # For the blocks are all padding + attn_weights = attn_weights.masked_fill(all_pad_mask, 0.0) + if torch.jit.is_scripting() or torch.jit.is_tracing(): pass elif random.random() < 0.001 and not self.training: @@ -1678,7 +1720,7 @@ class RelPositionMultiheadAttentionWeights(nn.Module): # (head, batch, time1, pos_dim) x (head, 1, pos_dim, seq_len2) -> (head, batch, time1, seq_len2) # [where seq_len2 represents relative position.] pos_scores = torch.matmul(p, pos_emb) - + if torch.jit.is_tracing(): (num_heads, batch_size, time1, n) = pos_scores.shape rows = torch.arange(start=time1 - 1, end=-1, step=-1) @@ -1743,8 +1785,10 @@ class SelfAttention(nn.Module): embed_dim: int, num_heads: int, value_head_dim: int, + block_size: int, ) -> None: super().__init__() + self.block_size = block_size self.in_proj = nn.Linear(embed_dim, num_heads * value_head_dim, bias=True) @@ -1766,27 +1810,45 @@ class SelfAttention(nn.Module): """ Args: x: input tensor, of shape (seq_len, batch_size, embed_dim) - attn_weights: a tensor of shape (num_heads, batch_size, seq_len, seq_len), - with seq_len being interpreted as (tgt_seq_len, src_seq_len). Expect - attn_weights.sum(dim=-1) == 1. + attn_weights: a tensor of attention weights, of shape + (hum_heads, batch_size * num_blocks, block_size, block_size * 3) + interpreted as (hum_heads, batch_size * num_blocks, tgt_seq_len, src_seq_len), + where num_blocks = (seq_len + block_size - 1) // block_size. + Expect attn_weights.sum(dim=-1) == 1. Returns: a tensor with the same shape as x. """ (seq_len, batch_size, embed_dim) = x.shape num_heads = attn_weights.shape[0] - assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len) + + # divide into blocks by unfold function + block_size = self.block_size + num_blocks = (seq_len + block_size - 1) // block_size + pad_len = num_blocks * block_size - seq_len + new_batch_size = batch_size * num_blocks + time1 = block_size # target length + time2 = 3 * block_size # source length + + assert attn_weights.shape == (num_heads, new_batch_size, time1, time2) x = self.in_proj(x) # (seq_len, batch_size, num_heads * value_head_dim) - x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3) - # now x: (num_heads, batch_size, seq_len, value_head_dim) - value_head_dim = x.shape[-1] + + # (time2, new_batch_size, channel) + x_blocks = unfold(x, pad_len, num_blocks, kernel=block_size * 3, stride=block_size, padding=block_size) + + x_blocks = x_blocks.reshape(time2, new_batch_size, num_heads, -1).permute(2, 1, 0, 3) + # now x: (num_heads, new_batch_size, time2, value_head_dim) + value_head_dim = x_blocks.shape[-1] # todo: see whether there is benefit in overriding matmul - x = torch.matmul(attn_weights, x) - # v: (num_heads, batch_size, seq_len, value_head_dim) + x = torch.matmul(attn_weights, x_blocks) + # v: (num_heads, new_batch_size, time1, value_head_dim) - x = x.permute(2, 1, 0, 3).contiguous().view( - seq_len, batch_size, num_heads * value_head_dim) + x = x.reshape(num_heads, batch_size, num_blocks, time1, value_head_dim) + x = x.permute(2, 3, 1, 0, 4).contiguous().view( + num_blocks * time1, batch_size, num_heads * value_head_dim) + + x = x[:seq_len] # (seq_len, batch_size, value_dim) # returned value is of shape (seq_len, batch_size, embed_dim), like the input. x = self.out_proj(x) @@ -1896,10 +1958,12 @@ class NonlinAttention(nn.Module): self, channels: int, hidden_channels: int, + block_size: int, ) -> None: super().__init__() self.hidden_channels = hidden_channels + self.block_size = block_size self.in_proj = nn.Linear(channels, hidden_channels * 3, bias=True) @@ -1942,7 +2006,11 @@ class NonlinAttention(nn.Module): """. Args: x: a Tensor of shape (seq_len, batch_size, num_channels) -attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len) + attn_weights: a tensor of attention weights, of shape + (hum_heads, batch_size * num_blocks, block_size, block_size * 3) + interpreted as (hum_heads, batch_size * num_blocks, tgt_seq_len, src_seq_len), + where num_blocks = (seq_len + block_size - 1) // block_size. + Expect attn_weights.sum(dim=-1) == 1. Returns: a Tensor with the same shape as x """ @@ -1965,13 +2033,31 @@ attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len) (seq_len, batch_size, embed_dim) = x.shape num_heads = attn_weights.shape[0] - assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len) - x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3) - # now x: (num_heads, batch_size, seq_len, head_dim) - x = torch.matmul(attn_weights, x) - # now x: (num_heads, batch_size, seq_len, head_dim) - x = x.permute(2, 1, 0, 3).reshape(seq_len, batch_size, -1) + # divide into blocks by unfold function + block_size = self.block_size + num_blocks = (seq_len + block_size - 1) // block_size + pad_len = num_blocks * block_size - seq_len + new_batch_size = batch_size * num_blocks + time1 = block_size # target length + time2 = 3 * block_size # source length + + assert attn_weights.shape == (num_heads, new_batch_size, time1, time2) + + # (time2, new_batch_size, channel) + x_blocks = unfold(x, pad_len, num_blocks, kernel=block_size * 3, stride=block_size, padding=block_size) + + x_blocks = x_blocks.reshape(time2, new_batch_size, num_heads, -1).permute(2, 1, 0, 3) + # now x: (num_heads, new_batch_size, time2, head_dim) + + x = torch.matmul(attn_weights, x_blocks) + # now x: (num_heads, new_batch_size, time1, head_dim) + + x = x.reshape(num_heads, batch_size, num_blocks, time1, -1) + x = x.permute(2, 3, 1, 0, 4).contiguous().view( + num_blocks * time1, batch_size, embed_dim) + + x = x[:seq_len] # (seq_len, batch_size, embed_dim) y = self.identity2(y) x = x * y @@ -2220,30 +2306,37 @@ class ScalarMultiply(nn.Module): def _test_zipformer_main(causal: bool = False): - batch_size = 5 - seq_len = 20 # Just make sure the forward pass runs. + from icefall.utils import make_pad_mask + c = Zipformer2( encoder_dim=(64, 96), encoder_unmasked_dim=(48, 64), num_heads=(4, 4), + downsampling_factor=(1, 2), + block_size=4, causal=causal, chunk_size=(4,) if causal else (-1,), left_context_frames=(64,) ) - batch_size = 5 - seq_len = 20 + batch_size = 2 + seq_len = 14 + # Just make sure the forward pass runs. - f = c( - torch.randn(seq_len, batch_size, 64), - torch.full((batch_size,), seq_len, dtype=torch.int64), - ) + x = torch.randn(seq_len, batch_size, 64) + lengths = torch.full((batch_size,), seq_len, dtype=torch.int64) + lengths[-1] = 1 + src_key_padding_mask = make_pad_mask(lengths) + f = c(x, lengths, src_key_padding_mask) f[0].sum().backward() c.eval() - f = c( - torch.randn(seq_len, batch_size, 64), - torch.full((batch_size,), seq_len, dtype=torch.int64), - ) + + x = torch.randn(seq_len, batch_size, 64) + lengths = torch.full((batch_size,), seq_len, dtype=torch.int64) + lengths[-1] = seq_len - 2 + src_key_padding_mask = make_pad_mask(lengths) + f = c(x, lengths, src_key_padding_mask) f # to remove flake8 warnings + print(f[0].sum()) if __name__ == "__main__": @@ -2251,4 +2344,4 @@ if __name__ == "__main__": torch.set_num_threads(1) torch.set_num_interop_threads(1) _test_zipformer_main(False) - _test_zipformer_main(True) + # _test_zipformer_main(True)