#!/usr/bin/env python3 # Copyright 2022-2023 Xiaomi Corp. (authors: Daniel Povey, # Zengwei Yao) # Copyright (c) 2024 Alibaba, Inc. and its affiliates. # # See ../../../../LICENSE for clarification regarding multiple authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import logging import math import random import warnings from typing import List, Optional, Tuple, Union import torch from torch import Tensor, nn from .scaling import \ Identity # more friendly to backward hooks than nn.Identity(), for diagnostic reasons. from .scaling import \ ScaledLinear # not as in other dirs.. just scales down initial parameter values. from .scaling import (ActivationDropoutAndLinear, BiasNorm, ChunkCausalDepthwiseConv1d, Dropout2, FloatLike, ScheduledFloat, limit_param_value, penalize_abs_values_gt, softmax) class Zipformer2EncoderLayer(nn.Module): """ Args: embed_dim: the number of expected features in the input (required). nhead: the number of heads in the multiheadattention models (required). feedforward_dim: the dimension of the feedforward network model (required). dropout: the dropout value (default=0.1). cnn_module_kernel (int): Kernel size of convolution module (default=31). Examples:: >>> encoder_layer = Zipformer2EncoderLayer(embed_dim=512, nhead=8) >>> src = torch.rand(10, 32, 512) >>> pos_emb = torch.rand(32, 19, 512) >>> out = encoder_layer(src, pos_emb) """ def __init__( self, embed_dim: int, pos_dim: int, num_heads: int, query_head_dim: int, pos_head_dim: int, value_head_dim: int, feedforward_dim: int, dropout: FloatLike = 0.1, cnn_module_kernel: int = 31, 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), const_attention_rate: FloatLike = ScheduledFloat((0.0, 0.25), (4000.0, 0.025), default=0), ff2_skip_rate: FloatLike = ScheduledFloat((0.0, 0.1), (4000.0, 0.01), (50000.0, 0.0)), ff3_skip_rate: FloatLike = ScheduledFloat((0.0, 0.1), (4000.0, 0.01), (50000.0, 0.0)), bypass_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.02), default=0), ) -> None: super(Zipformer2EncoderLayer, self).__init__() self.embed_dim = embed_dim # self.bypass implements layer skipping as well as bypass; see its default values. self.bypass = BypassModule( embed_dim, skip_rate=bypass_skip_rate, straight_through_rate=0) # bypass_mid is bypass used in the middle of the layer. self.bypass_mid = BypassModule(embed_dim, straight_through_rate=0) # skip probability for dynamic modules (meaning: anything but feedforward). self.attention_skip_rate = copy.deepcopy(attention_skip_rate) # an additional skip probability that applies to ConvModule to stop it from # contributing too much early on. self.conv_skip_rate = copy.deepcopy(conv_skip_rate) # ff2_skip_rate is to prevent the ff2 module from having output that's too big # compared to its residual. self.ff2_skip_rate = copy.deepcopy(ff2_skip_rate) self.ff3_skip_rate = copy.deepcopy(ff3_skip_rate) self.const_attention_rate = copy.deepcopy(const_attention_rate) 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, ) self.self_attn1 = SelfAttention(embed_dim, num_heads, value_head_dim) self.self_attn2 = SelfAttention(embed_dim, num_heads, value_head_dim) self.feed_forward1 = FeedforwardModule(embed_dim, (feedforward_dim * 3) // 4, dropout) self.feed_forward2 = FeedforwardModule(embed_dim, feedforward_dim, dropout) self.feed_forward3 = FeedforwardModule(embed_dim, (feedforward_dim * 5) // 4, dropout) self.nonlin_attention = NonlinAttention( embed_dim, hidden_channels=3 * embed_dim // 4) self.conv_module1 = ConvolutionModule( embed_dim, cnn_module_kernel, causal=causal) self.conv_module2 = ConvolutionModule( embed_dim, cnn_module_kernel, causal=causal) # TODO: remove it self.bypass_scale = nn.Parameter(torch.full((embed_dim, ), 0.5)) self.norm = BiasNorm(embed_dim) self.balancer1 = Identity() self.balancer_na = Identity() self.balancer_ff2 = Identity() self.balancer_ff3 = Identity() self.whiten = Identity() self.balancer2 = Identity() def get_sequence_dropout_mask(self, x: Tensor, dropout_rate: float) -> Optional[Tensor]: if (dropout_rate == 0.0 or not self.training or torch.jit.is_scripting() or torch.jit.is_tracing()): return None batch_size = x.shape[1] mask = (torch.rand(batch_size, 1, device=x.device) > dropout_rate).to( x.dtype) return mask def sequence_dropout(self, x: Tensor, dropout_rate: float) -> Tensor: """ Apply sequence-level dropout to x. x shape: (seq_len, batch_size, embed_dim) """ dropout_mask = self.get_sequence_dropout_mask(x, dropout_rate) if dropout_mask is None: return x else: return x * dropout_mask def forward( self, src: Tensor, pos_emb: Tensor, chunk_size: int = -1, attn_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, ) -> Tensor: """ Pass the input through the encoder layer. Args: src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim). pos_emb: (1, 2*seq_len-1, pos_emb_dim) or (batch_size, 2*seq_len-1, pos_emb_dim) chunk_size: the number of frames per chunk, of >= 0; if -1, no chunking. 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 which has the same shape as src """ src_orig = src # dropout rate for non-feedforward submodules if torch.jit.is_scripting() or torch.jit.is_tracing(): attention_skip_rate = 0.0 else: attention_skip_rate = ( float(self.attention_skip_rate) if self.training else 0.0) # attn_weights: (num_heads, batch_size, seq_len, seq_len) attn_weights = self.self_attn_weights( src, pos_emb=pos_emb, attn_mask=attn_mask, key_padding_mask=src_key_padding_mask, ) src = src + self.feed_forward1(src) self_attn_dropout_mask = self.get_sequence_dropout_mask( src, attention_skip_rate) selected_attn_weights = attn_weights[0:1] if torch.jit.is_scripting() or torch.jit.is_tracing(): pass elif self.training and random.random() < float( self.const_attention_rate): # Make attention weights constant. The intention is to # encourage these modules to do something similar to an # averaging-over-time operation. # only need the mask, can just use the 1st one and expand later selected_attn_weights = selected_attn_weights[0:1] selected_attn_weights = (selected_attn_weights > 0.0).to( selected_attn_weights.dtype) selected_attn_weights = selected_attn_weights * ( 1.0 / selected_attn_weights.sum(dim=-1, keepdim=True)) na = self.balancer_na( self.nonlin_attention(src, selected_attn_weights)) src = src + ( na if self_attn_dropout_mask is None else na * self_attn_dropout_mask) self_attn = self.self_attn1(src, attn_weights) src = src + ( self_attn if self_attn_dropout_mask is None else self_attn * self_attn_dropout_mask) if torch.jit.is_scripting() or torch.jit.is_tracing(): conv_skip_rate = 0.0 else: conv_skip_rate = float( self.conv_skip_rate) if self.training else 0.0 src = src + self.sequence_dropout( self.conv_module1( src, chunk_size=chunk_size, src_key_padding_mask=src_key_padding_mask), conv_skip_rate, ) if torch.jit.is_scripting() or torch.jit.is_tracing(): ff2_skip_rate = 0.0 else: ff2_skip_rate = float(self.ff2_skip_rate) if self.training else 0.0 src = src + self.sequence_dropout( self.balancer_ff2(self.feed_forward2(src)), ff2_skip_rate) # bypass in the middle of the layer. src = self.bypass_mid(src_orig, src) self_attn = self.self_attn2(src, attn_weights) src = src + ( self_attn if self_attn_dropout_mask is None else self_attn * self_attn_dropout_mask) if torch.jit.is_scripting() or torch.jit.is_tracing(): conv_skip_rate = 0.0 else: conv_skip_rate = float( self.conv_skip_rate) if self.training else 0.0 src = src + self.sequence_dropout( self.conv_module2( src, chunk_size=chunk_size, src_key_padding_mask=src_key_padding_mask), conv_skip_rate, ) if torch.jit.is_scripting() or torch.jit.is_tracing(): ff3_skip_rate = 0.0 else: ff3_skip_rate = float(self.ff3_skip_rate) if self.training else 0.0 src = src + self.sequence_dropout( self.balancer_ff3(self.feed_forward3(src)), ff3_skip_rate) src = self.balancer1(src) src = self.norm(src) src = self.bypass(src_orig, src) src = self.balancer2(src) src = self.whiten(src) return src class BypassModule(nn.Module): """ An nn.Module that implements a learnable bypass scale, and also randomized per-sequence layer-skipping. The bypass is limited during early stages of training to be close to "straight-through", i.e. to not do the bypass operation much initially, in order to force all the modules to learn something. """ def __init__( self, embed_dim: int, skip_rate: FloatLike = 0.0, straight_through_rate: FloatLike = 0.0, scale_min: FloatLike = ScheduledFloat((0.0, 0.9), (20000.0, 0.2), default=0), scale_max: FloatLike = 1.0, ): super().__init__() self.bypass_scale = nn.Parameter(torch.full((embed_dim, ), 0.5)) self.skip_rate = copy.deepcopy(skip_rate) self.straight_through_rate = copy.deepcopy(straight_through_rate) self.scale_min = copy.deepcopy(scale_min) self.scale_max = copy.deepcopy(scale_max) def _get_bypass_scale(self, batch_size: int): # returns bypass-scale of shape (num_channels,), # or (batch_size, num_channels,). This is actually the # scale on the non-residual term, so 0 corresponds to bypassing # this module. if torch.jit.is_scripting() or torch.jit.is_tracing( ) or not self.training: return self.bypass_scale else: ans = limit_param_value( self.bypass_scale, min=float(self.scale_min), max=float(self.scale_max)) skip_rate = float(self.skip_rate) if skip_rate != 0.0: mask = torch.rand( (batch_size, 1), device=ans.device) > skip_rate ans = ans * mask # now ans is of shape (batch_size, num_channels), and is zero for sequences # on which we have randomly chosen to do layer-skipping. straight_through_rate = float(self.straight_through_rate) if straight_through_rate != 0.0: _rand_tensor = torch.rand((batch_size, 1), device=ans.device) mask = (_rand_tensor < straight_through_rate) ans = torch.maximum(ans, mask.to(ans.dtype)) return ans def forward(self, src_orig: Tensor, src: Tensor): """ Args: src_orig and src are both of shape (seq_len, batch_size, num_channels) Returns: something with the same shape as src and src_orig """ # bypass_scale = self._get_bypass_scale(src.shape[1]) bypass_scale = self._get_bypass_scale(src.shape[-2]) return src_orig + (src - src_orig) * bypass_scale class SimpleDownsample(torch.nn.Module): """ Does downsampling with attention, by weighted sum, and a projection.. """ def __init__(self, channels: int, downsample: int, dropout: FloatLike): 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) -> 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 # 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): 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 class CompactRelPositionalEncoding(torch.nn.Module): """ Relative positional encoding module. This version is "compact" meaning it is able to encode the important information about the relative position in a relatively small number of dimensions. The goal is to make it so that small differences between large relative offsets (e.g. 1000 vs. 1001) make very little difference to the embedding. Such differences were potentially important when encoding absolute position, but not important when encoding relative position because there is now no need to compare two large offsets with each other. Our embedding works by projecting the interval [-infinity,infinity] to a finite interval using the atan() function, before doing the Fourier transform of that fixed interval. The atan() function would compress the "long tails" too small, making it hard to distinguish between different magnitudes of large offsets, so we use a logarithmic function to compress large offsets to a smaller range before applying atan(). Scalings are chosen in such a way that the embedding can clearly distinguish individual offsets as long as they are quite close to the origin, e.g. abs(offset) <= about sqrt(embedding_dim) Args: embed_dim: Embedding dimension. dropout_rate: Dropout rate. max_len: Maximum input length: just a heuristic for initialization. length_factor: a heuristic scale (should be >= 1.0) which, if larger, gives less weight to small differences of offset near the origin. """ def __init__( self, embed_dim: int, dropout_rate: FloatLike, max_len: int = 1000, length_factor: float = 1.0, ) -> None: """Construct a CompactRelPositionalEncoding object.""" super(CompactRelPositionalEncoding, self).__init__() self.embed_dim = embed_dim assert embed_dim % 2 == 0, embed_dim self.dropout = Dropout2(dropout_rate) self.pe = None assert length_factor >= 1.0, length_factor 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: """Reset the positional encodings.""" T = x.size(0) + left_context_len if self.pe is not None: # self.pe contains both positive and negative parts # the length of self.pe is 2 * input_len - 1 if self.pe.size(0) >= T * 2 - 1: self.pe = self.pe.to(dtype=x.dtype, device=x.device) return # if T == 4, x would contain [ -3, -2, 1, 0, 1, 2, 3 ] x = torch.arange( -(T - 1), T, device=x.device).to(torch.float32).unsqueeze(1) freqs = 1 + torch.arange(self.embed_dim // 2, device=x.device) # `compression_length` this is arbitrary/heuristic, if it is larger we have more resolution # for small time offsets but less resolution for large time offsets. compression_length = self.embed_dim**0.5 # x_compressed, like X, goes from -infinity to infinity as T goes from -infinity to infinity; # but it does so more slowly than T for large absolute values of T. # The formula is chosen so that d(x_compressed )/dx is 1 around x == 0, which # is important. _tmp_tensor = ((x.abs() + compression_length).log() - math.log(compression_length)) x_compressed = (compression_length * x.sign() * _tmp_tensor) # if self.length_factor == 1.0, then length_scale is chosen so that the # FFT can exactly separate points close to the origin (T == 0). So this # part of the formulation is not really heuristic. # But empirically, for ASR at least, length_factor > 1.0 seems to work better. length_scale = self.length_factor * self.embed_dim / (2.0 * math.pi) # note for machine implementations: if atan is not available, we can use: # x.sign() * ((1 / (x.abs() + 1)) - 1) * (-math.pi/2) # check on wolframalpha.com: plot(sign(x) * (1 / ( abs(x) + 1) - 1 ) * -pi/2 , atan(x)) x_atan = (x_compressed / length_scale).atan() # results between -pi and pi cosines = (x_atan * freqs).cos() sines = (x_atan * freqs).sin() pe = torch.zeros(x.shape[0], self.embed_dim, device=x.device) pe[:, 0::2] = cosines pe[:, 1::2] = sines pe[:, -1] = 1.0 # for bias. self.pe = pe.to(dtype=x.dtype) def forward(self, x: Tensor, left_context_len: int = 0) -> Tensor: """Create positional encoding. Args: x (Tensor): Input tensor (time, batch, `*`). left_context_len: (int): Length of cached left context. Returns: positional embedding, of shape (batch, left_context_len + 2*time-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) pos_emb = self.pe[self.pe.size(0) // 2 - x_size_left + 1:self.pe.size(0) // 2 # noqa E203 + x.size(0), :, ] pos_emb = pos_emb.unsqueeze(0) return self.dropout(pos_emb) class RelPositionMultiheadAttentionWeights(nn.Module): r"""Module that computes multi-head attention weights with relative position encoding. Various other modules consume the resulting attention weights: see, for example, the SimpleAttention module which allows you to compute conventional attention. This is a quite heavily modified from: "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context", we have to write up the differences. Args: embed_dim: number of channels at the input to this module, e.g. 256 pos_dim: dimension of the positional encoding vectors, e.g. 128. num_heads: number of heads to compute weights for, e.g. 8 query_head_dim: dimension of the query (and key), per head. e.g. 24. pos_head_dim: dimension of the projected positional encoding per head, e.g. 4. dropout: dropout probability for attn_output_weights. Default: 0.0. pos_emb_skip_rate: probability for skipping the pos_emb part of the scores on any given call to forward(), in training time. """ def __init__( self, embed_dim: int, pos_dim: int, num_heads: int, query_head_dim: int, pos_head_dim: int, dropout: float = 0.0, pos_emb_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.0)), ) -> None: super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.query_head_dim = query_head_dim self.pos_head_dim = pos_head_dim 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. key_head_dim = query_head_dim in_proj_dim = (query_head_dim + key_head_dim + pos_head_dim) * num_heads # the initial_scale is supposed to take over the "scaling" factor of # head_dim ** -0.5 that has been used in previous forms of attention, # dividing it between the query and key. Note: this module is intended # to be used with the ScaledAdam optimizer; with most other optimizers, # it would be necessary to apply the scaling factor in the forward function. self.in_proj = ScaledLinear( embed_dim, in_proj_dim, bias=True, initial_scale=query_head_dim**-0.25) self.whiten_keys = Identity() self.balance_keys = Identity() # linear transformation for positional encoding. self.linear_pos = ScaledLinear( pos_dim, num_heads * pos_head_dim, bias=False, initial_scale=0.05) # the following are for diagnostics only, see --print-diagnostics option self.copy_pos_query = Identity() self.copy_query = Identity() def forward( self, x: Tensor, pos_emb: Tensor, key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None, ) -> Tensor: 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) 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). """ x = self.in_proj(x) query_head_dim = self.query_head_dim pos_head_dim = self.pos_head_dim num_heads = self.num_heads seq_len, batch_size, _ = x.shape query_dim = query_head_dim * num_heads # self-attention q = x[..., 0:query_dim] k = x[..., query_dim:2 * query_dim] # p is the position-encoding query p = x[..., 2 * query_dim:] assert p.shape[-1] == num_heads * pos_head_dim, (p.shape[-1], num_heads, pos_head_dim) q = self.copy_query(q) # for diagnostics only, does nothing. 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) # 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) # print(f"MHSAW {q.shape} {k.shape}") attn_scores = torch.matmul(q, k) use_pos_scores = False if torch.jit.is_scripting() or torch.jit.is_tracing(): # We can't put random.random() in the same line use_pos_scores = True elif not self.training or random.random() >= float( self.pos_emb_skip_rate): use_pos_scores = True 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) # (head, batch, time1, pos_dim) x (head, 1, pos_dim, seq_len2) -> (head, batch, time1, seq_len2) # [where seq_len2 represents relative position.] # print(f"MHSAW pos {p.shape} {pos_emb.shape}") pos_scores = torch.matmul(p, 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), ) # print(attn_scores.shape, pos_scores.shape) if self.training: attn_scores = attn_scores + pos_scores else: # inplace operator important attn_scores.add_(pos_scores) # attn_scores = attn_scores + pos_scores if torch.jit.is_scripting() or torch.jit.is_tracing(): pass elif self.training and random.random() < 0.1: # This is a harder way of limiting the attention scores to not be # too large. It incurs a penalty if any of them has an absolute # value greater than 50.0. this should be outside the normal range # of the attention scores. We use this mechanism instead of, say, # something added to the loss function involving the entropy, # because once the entropy gets very small gradients through the # softmax can become very small, and we'd get zero derivatives. The # choices of 1.0e-04 as the scale on the penalty makes this # mechanism vulnerable to the absolute scale of the loss function, # but we view this as a failsafe to avoid "implausible" parameter # values rather than a regularization method that should be active # under normal circumstances. attn_scores = penalize_abs_values_gt( attn_scores, limit=25.0, penalty=1.0e-04, name=self.name) assert attn_scores.shape == (num_heads, batch_size, seq_len, seq_len) if attn_mask is not None: 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) # is exactly zero, for reasons related to const_attention_rate, it # 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, ) # 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 # automatic mixed precision mode (amp / autocast), by only storing the # half-precision output for backprop purposes. attn_weights = softmax(attn_scores, dim=-1) attn_weights = nn.functional.dropout( attn_weights, p=self.dropout, training=self.training) return attn_weights class SelfAttention(nn.Module): """ The simplest possible attention module. This one works with already-computed attention weights, e.g. as computed by RelPositionMultiheadAttentionWeights. Args: embed_dim: the input and output embedding dimension num_heads: the number of attention heads value_head_dim: the value dimension per head """ def __init__( self, embed_dim: int, num_heads: int, value_head_dim: int, ) -> None: super().__init__() self.in_proj = nn.Linear( embed_dim, num_heads * value_head_dim, bias=True) self.out_proj = ScaledLinear( num_heads * value_head_dim, embed_dim, bias=True, initial_scale=0.05) self.whiten = Identity() def forward( self, x: Tensor, attn_weights: Tensor, ) -> Tensor: """ 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. 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) 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] # todo: see whether there is benefit in overriding matmul # print(f"SelfAttetion pos {attn_weights.shape} {x.shape}") x = torch.matmul(attn_weights, x) # v: (num_heads, batch_size, seq_len, value_head_dim) x = ( x.permute(2, 1, 0, 3).contiguous().view(seq_len, batch_size, num_heads * value_head_dim)) # returned value is of shape (seq_len, batch_size, embed_dim), like the input. x = self.out_proj(x) x = self.whiten(x) return x class FeedforwardModule(nn.Module): """Feedforward module in Zipformer2 model.""" def __init__(self, embed_dim: int, feedforward_dim: int, dropout: FloatLike): super(FeedforwardModule, self).__init__() self.in_proj = nn.Linear(embed_dim, feedforward_dim) self.hidden_balancer = Identity() # shared_dim=0 means we share the dropout mask along the time axis self.out_proj = ActivationDropoutAndLinear( feedforward_dim, embed_dim, activation='SwooshL', dropout_p=dropout, dropout_shared_dim=0, bias=True, initial_scale=0.1, ) self.out_whiten = Identity() def forward(self, x: Tensor): x = self.in_proj(x) x = self.hidden_balancer(x) # out_proj contains SwooshL activation, then dropout, then linear. x = self.out_proj(x) x = self.out_whiten(x) return x class NonlinAttention(nn.Module): """This is like the ConvolutionModule, but refactored so that we use multiplication by attention weights (borrowed from the attention module) in place of actual convolution. We also took out the second nonlinearity, the one after the attention mechanism. Args: channels (int): The number of channels of conv layers. """ def __init__( self, channels: int, hidden_channels: int, ) -> None: super().__init__() self.hidden_channels = hidden_channels self.in_proj = nn.Linear(channels, hidden_channels * 3, bias=True) self.balancer = Identity() self.tanh = nn.Tanh() self.identity1 = Identity() # for diagnostics. self.identity2 = Identity() # for diagnostics. self.identity3 = Identity() # for diagnostics. self.out_proj = ScaledLinear( hidden_channels, channels, bias=True, initial_scale=0.05) self.whiten1 = Identity() self.whiten2 = Identity() def forward( self, x: Tensor, attn_weights: Tensor, ) -> Tensor: """. 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) Returns: a Tensor with the same shape as x """ x = self.in_proj(x) (seq_len, batch_size, _) = x.shape hidden_channels = self.hidden_channels s, x, y = x.chunk(3, dim=2) # s will go through tanh. s = self.balancer(s) s = self.tanh(s) s = s.unsqueeze(-1).reshape(seq_len, batch_size, hidden_channels) x = self.whiten1(x) x = x * s x = self.identity1(x) # diagnostics only, it's the identity. (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) # print(f"nonlinattion {attn_weights.shape} {x.shape}") 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) y = self.identity2(y) x = x * y x = self.identity3(x) x = self.out_proj(x) x = self.whiten2(x) return x class ConvolutionModule(nn.Module): """ConvolutionModule in Zipformer2 model. Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/zipformer/convolution.py Args: channels (int): The number of channels of conv layers. kernel_size (int): Kernel size of conv layers. bias (bool): Whether to use bias in conv layers (default=True). """ def __init__( self, channels: int, kernel_size: int, causal: bool, ) -> None: """Construct a ConvolutionModule object.""" super(ConvolutionModule, self).__init__() # kernerl_size should be a odd number for 'SAME' padding assert (kernel_size - 1) % 2 == 0 bottleneck_dim = channels self.causal = causal self.in_proj = nn.Linear( channels, 2 * bottleneck_dim, ) # the gradients on in_proj are a little noisy, likely to do with the # sigmoid in glu. self.balancer1 = Identity() self.activation1 = Identity() # for diagnostics self.sigmoid = nn.Sigmoid() self.activation2 = Identity() # for diagnostics assert kernel_size % 2 == 1 self.depthwise_conv = ( ChunkCausalDepthwiseConv1d( channels=bottleneck_dim, kernel_size=kernel_size) if causal else nn.Conv1d( in_channels=bottleneck_dim, out_channels=bottleneck_dim, groups=bottleneck_dim, kernel_size=kernel_size, padding=kernel_size // 2, )) self.balancer2 = Identity() self.whiten = Identity() self.out_proj = ActivationDropoutAndLinear( bottleneck_dim, channels, activation='SwooshR', dropout_p=0.0, initial_scale=0.05, ) def forward( self, x: Tensor, src_key_padding_mask: Optional[Tensor] = None, chunk_size: int = -1, ) -> Tensor: """Compute convolution module. Args: x: Input tensor (#time, batch, channels). src_key_padding_mask: the mask for the src keys per batch (optional): (batch, #time), contains True in masked positions. Returns: Tensor: Output tensor (#time, batch, channels). """ x = self.in_proj(x) # (time, batch, 2*channels) x, s = x.chunk(2, dim=2) s = self.balancer1(s) s = self.sigmoid(s) x = self.activation1(x) # identity. x = x * s x = self.activation2(x) # identity # (time, batch, channels) # exchange the temporal dimension and the feature dimension x = x.permute(1, 2, 0) # (#batch, channels, time). if src_key_padding_mask is not None: x = x.masked_fill( src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0) if (not torch.jit.is_scripting() and not torch.jit.is_tracing() and chunk_size >= 0): # Not support exporting a model for simulated streaming decoding assert ( self.causal ), 'Must initialize model with causal=True if you use chunk_size' x = self.depthwise_conv(x, chunk_size=chunk_size) else: # with record_function("depthwise_conv"): x = self.depthwise_conv(x) # pass x = self.balancer2(x) x = x.permute(2, 0, 1) # (time, batch, channels) x = self.whiten(x) # (time, batch, channels) x = self.out_proj(x) # (time, batch, channels) return x if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) torch.set_num_threads(1) torch.set_num_interop_threads(1)