mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-12-11 06:55:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove squeeze_excite from Conv2dSubsampling.

Browse Source
This commit is contained in:
Daniel Povey 2022-11-26 14:04:41 +08:00
parent 281b54e7bf
commit 4058d56c0d
1 changed files with 0 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
View File

@ -1606,48 +1606,6 @@ class ConvolutionModule(nn.Module):
return x
class SqueezeExcite1d(nn.Module):
def __init__(self,
channels: int,
bottleneck_channels: int):
super().__init__()
self.to_bottleneck_proj = LinearWithAuxLoss(channels,
bottleneck_channels)
self.bottleneck_activation = TanSwish()
self.from_bottleneck_proj = nn.Linear(bottleneck_channels,
channels)
self.balancer = ActivationBalancer(
channels, channel_dim=-1,
min_abs=0.05,
max_abs=ScheduledFloat((0.0, 0.2),
(4000.0, 2.0),
(10000.0, 10.0),
default=1.0),
max_factor=0.02,
min_prob=0.1,
)
self.activation = nn.Sigmoid()
def forward(self, x: Tensor):
"""
x: a Tensor of shape (batch_size, T, channels).
Returns: something with the same shape as x.
"""
# project before mean, needed for LinearWithAuxLoss (or, at least, better)
bottleneck = self.to_bottleneck_proj(x)
# would replace this mean with cumsum for a causal model.
bottleneck = bottleneck.mean(dim=1, keepdim=True)
bottleneck = self.bottleneck_activation(bottleneck)
scale = self.from_bottleneck_proj(bottleneck)
scale = self.balancer(scale)
scale = self.activation(scale)
return x * scale
class Conv2dSubsampling(nn.Module):
"""Convolutional 2D subsampling (to 1/2 length).
@ -1718,8 +1676,6 @@ class Conv2dSubsampling(nn.Module):
)
out_height = (((in_channels - 1) // 2) - 1) // 2
self.squeeze_excite = SqueezeExcite1d(out_height * layer3_channels,
bottleneck_channels)
self.out = LinearWithAuxLoss(out_height * layer3_channels, out_channels,
aux_grad_scale=ScheduledFloat((0.0, 0.2), (1000.0, 0.01)))
@ -1745,7 +1701,6 @@ class Conv2dSubsampling(nn.Module):
x = x.transpose(1, 2).reshape(b, t, c * f)
# now x: (N, ((T-1)//2 - 1))//2, out_height * layer3_channels))
x = self.squeeze_excite(x)
x = self.out(x)
# Now x is of shape (N, ((T-1)//2 - 1))//2, odim)
x = self.dropout(x)