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

Have 2 not 3 groups, but give 1st group a smaller dropout prob than the 2nd.

Browse Source
This commit is contained in:
Daniel Povey 2023-03-30 16:38:41 +08:00
parent 6e058b9ebd
commit e64ec396bd
1 changed files with 11 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -297,30 +297,25 @@ class Zipformer2(EncoderInterface):
num_frames_max = (num_frames0 + max_downsampling_factor - 1)
# we divide the dropped-out feature dimensions into two equal groups;
# the first group is dropped out with this probability, the second
# group is dropped out with about twice this probability.
feature_mask_dropout_prob = 0.1
# the first group is dropped out with probability 0.05, the second
# with probability approximately (0.2 + 0.05).
feature_mask_dropout_prob1 = 0.05
feature_mask_dropout_prob2 = 0.2
# frame_mask_max1 shape: (num_frames_max, batch_size, 1)
frame_mask_max1 = (torch.rand(num_frames_max, batch_size, 1,
device=x.device) >
feature_mask_dropout_prob).to(x.dtype)
feature_mask_dropout_prob1).to(x.dtype)
# frame_mask_max2 has additional frames masked, about twice the number.
frame_mask_max2 = torch.logical_and(frame_mask_max1,
(torch.rand(num_frames_max, batch_size, 1,
device=x.device) >
feature_mask_dropout_prob).to(x.dtype))
feature_mask_dropout_prob2).to(x.dtype))
# frame_mask_max3 has additional frames masked, about 3 times the number.
frame_mask_max3 = torch.logical_and(frame_mask_max2,
(torch.rand(num_frames_max, batch_size, 1,
device=x.device) >
feature_mask_dropout_prob).to(x.dtype))
# dim: (num_frames_max, batch_size, 3)
frame_mask_max = torch.cat((frame_mask_max1, frame_mask_max2, frame_mask_max3), dim=-1)
frame_mask_max = torch.cat((frame_mask_max1, frame_mask_max2), dim=-1)
feature_masks = []
for i in range(num_encoders):
@ -328,19 +323,17 @@ class Zipformer2(EncoderInterface):
upsample_factor = (max_downsampling_factor // ds)
frame_mask = (frame_mask_max.unsqueeze(1).expand(num_frames_max, upsample_factor,
batch_size, 3)
.reshape(num_frames_max * upsample_factor, batch_size, 3))
batch_size, 2)
.reshape(num_frames_max * upsample_factor, batch_size, 2))
num_frames = (num_frames0 + ds - 1) // ds
frame_mask = frame_mask[:num_frames]
channels = self.encoder_dim[i]
feature_mask = torch.ones(num_frames, batch_size, channels,
dtype=x.dtype, device=x.device)
u1 = self.encoder_unmasked_dim[i]
u2 = u1 + (channels - u1) // 3
u3 = u1 + 2 * (channels - u1) // 3
u2 = u1 + (channels - u1) // 2
feature_mask[:, :, u1:u2] *= frame_mask[..., 0:1]
feature_mask[:, :, u2:u3] *= frame_mask[..., 1:2]
feature_mask[:, :, u3:channels] *= frame_mask[..., 2:3]
feature_mask[:, :, u2:] *= frame_mask[..., 1:2]
feature_masks.append(feature_mask)
return feature_masks