From c533c30442cc07d6382a34444946d5f1ddb38648 Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Sat, 31 Dec 2022 15:10:52 +0800 Subject: [PATCH 1/4] Increase final conv_skip_rate from 0.0 to 0.01 --- .../ASR/pruned_transducer_stateless7/zipformer.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py index f96ef34a9..4e8deb88f 100644 --- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py @@ -403,8 +403,9 @@ class ZipformerEncoderLayer(nn.Module): # to work correctly. layer_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.05), default=0), 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), + conv_skip_rate: FloatLike = ScheduledFloat((0.0, 0.2), (4000.0, 0.05), (16000, 0.01), default=0), const_attention_rate: FloatLike = ScheduledFloat((0.0, 0.25), (4000.0, 0.025), default=0), + ff2_skip_rate: FloatLike = 0.01, bypass_min: FloatLike = ScheduledFloat((0.0, 0.75), (20000.0, 0.2), default=0), bypass_max: FloatLike = 1.0, ) -> None: @@ -418,6 +419,9 @@ class ZipformerEncoderLayer(nn.Module): # 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) # min and max for self.bypass_scale, applied with probability 0.5 to avoid grads # ever becoming zero. @@ -608,7 +612,8 @@ class ZipformerEncoderLayer(nn.Module): if torch.jit.is_scripting() or random.random() >= float(self.conv_skip_rate): src = src + self.conv_module(src, src_key_padding_mask=src_key_padding_mask) - src = src + self.balancer_ff2(self.feed_forward2(src)) + if torch.jit.is_scripting() or random.random() >= float(self.ff2_skip_rate): + src = src + self.balancer_ff2(self.feed_forward2(src)) src = self.balancer1(src) src = self.norm(src) From d48b2ccb45753711a89949e8fac9e19ee41ecf45 Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Sat, 31 Dec 2022 17:10:31 +0800 Subject: [PATCH 2/4] Reduce kernel size of convnext2 from 7 to 5. --- .../ASR/pruned_transducer_stateless7/zipformer.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py index 4e8deb88f..5a0d22d86 100644 --- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py @@ -1704,10 +1704,10 @@ class ConvNeXt(nn.Module): def __init__(self, channels: int, hidden_ratio: int = 3, + kernel_size: Tuple[int, int] = (7, 7), layerdrop_rate: FloatLike = None): super().__init__() - kernel_size = 7 - pad = (kernel_size - 1) // 2 + padding = ((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2) hidden_channels = channels * hidden_ratio if layerdrop_rate is None: layerdrop_rate = ScheduledFloat((0.0, 0.2), (20000.0, 0.015)) @@ -1717,8 +1717,8 @@ class ConvNeXt(nn.Module): in_channels=channels, out_channels=channels, groups=channels, - kernel_size=7, - padding=(3, 3)) + kernel_size=kernel_size, + padding=padding) self.pointwise_conv1 = nn.Conv2d( in_channels=channels, @@ -1869,9 +1869,9 @@ class Conv2dSubsampling(nn.Module): SwooshR(), ) - self.convnext2 = nn.Sequential(ConvNeXt(layer3_channels), - ConvNeXt(layer3_channels), - ConvNeXt(layer3_channels)) + self.convnext2 = nn.Sequential(ConvNeXt(layer3_channels, kernel_size=(5, 5)), + ConvNeXt(layer3_channels, kernel_size=(5, 5)), + ConvNeXt(layer3_channels, kernel_size=(5, 5))) out_width = (((in_channels - 1) // 2) - 1) // 2 From d0ae60400eb8b101327bcf58d00040bbef11bfa1 Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Sat, 31 Dec 2022 17:19:02 +0800 Subject: [PATCH 3/4] Decrease convnext1 kernel size from 7x7 to 5x5 --- egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py index 5a0d22d86..4a946db5a 100644 --- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py @@ -1847,8 +1847,8 @@ class Conv2dSubsampling(nn.Module): SwooshR(), ) - self.convnext1 = nn.Sequential(ConvNeXt(layer2_channels), - ConvNeXt(layer2_channels), + self.convnext1 = nn.Sequential(ConvNeXt(layer2_channels, kernel_size=(5, 5)), + ConvNeXt(layer2_channels, kernel_size=(5, 5)), BasicNorm(layer2_channels, channel_dim=1)) From 4a4d12c9944b18d9e81e5af22c1dc835deab8586 Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Sat, 31 Dec 2022 21:52:11 +0800 Subject: [PATCH 4/4] Revert kernel size of convnext2 from 5x5 to 7x7 --- .../ASR/pruned_transducer_stateless7/zipformer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py index 4a946db5a..4e67ed35e 100644 --- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py @@ -1869,9 +1869,9 @@ class Conv2dSubsampling(nn.Module): SwooshR(), ) - self.convnext2 = nn.Sequential(ConvNeXt(layer3_channels, kernel_size=(5, 5)), - ConvNeXt(layer3_channels, kernel_size=(5, 5)), - ConvNeXt(layer3_channels, kernel_size=(5, 5))) + self.convnext2 = nn.Sequential(ConvNeXt(layer3_channels, kernel_size=(7, 7)), + ConvNeXt(layer3_channels, kernel_size=(7, 7)), + ConvNeXt(layer3_channels, kernel_size=(7, 7))) out_width = (((in_channels - 1) // 2) - 1) // 2