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

Sometimes mask more frames.

Browse Source
This commit is contained in:
Daniel Povey 2023-03-29 13:08:52 +08:00
parent 4e36656cef
commit bb8cbd7598
1 changed files with 18 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -296,27 +296,38 @@ class Zipformer2(EncoderInterface):
num_frames_max = (num_frames0 + max_downsampling_factor - 1) num_frames_max = (num_frames0 + max_downsampling_factor - 1)
feature_mask_dropout_prob = 0.15 feature_mask_dropout_prob = 0.125
# frame_mask_max shape: (num_frames_max, batch_size, 1) # frame_mask_max1 shape: (num_frames_max, batch_size, 1)
frame_mask_max = (torch.rand(num_frames_max, batch_size, 1, frame_mask_max1 = (torch.rand(num_frames_max, batch_size, 1,
device=x.device) > device=x.device) >
feature_mask_dropout_prob).to(x.dtype) feature_mask_dropout_prob).to(x.dtype)
# frame_mask_max2 has additional frames masked, about twice the number.
frame_mask_max2 = torch.logical_or(frame_mask_max1,
(torch.rand(num_frames_max, batch_size, 1,
device=x.device) >
feature_mask_dropout_prob).to(x.dtype))
# dim: (num_frames_max, batch_size, 2)
frame_mask_max = torch.cat((frame_mask_max1, frame_mask_max2), dim=-1)
feature_masks = [] feature_masks = []
for i in range(num_encoders): for i in range(num_encoders):
ds = self.downsampling_factor[i] ds = self.downsampling_factor[i]
upsample_factor = (max_downsampling_factor // ds) upsample_factor = (max_downsampling_factor // ds)
frame_mask = (frame_mask_max.unsqueeze(1).expand(num_frames_max, upsample_factor, frame_mask = (frame_mask_max.unsqueeze(1).expand(num_frames_max, upsample_factor,
batch_size, 1) batch_size, 2)
.reshape(num_frames_max * upsample_factor, batch_size, 1)) .reshape(num_frames_max * upsample_factor, batch_size, 2))
num_frames = (num_frames0 + ds - 1) // ds num_frames = (num_frames0 + ds - 1) // ds
frame_mask = frame_mask[:num_frames] frame_mask = frame_mask[:num_frames]
feature_mask = torch.ones(num_frames, batch_size, self.encoder_dim[i], feature_mask = torch.ones(num_frames, batch_size, self.encoder_dim[i],
dtype=x.dtype, device=x.device) dtype=x.dtype, device=x.device)
u = self.encoder_unmasked_dim[i] u1 = self.encoder_unmasked_dim[i]
feature_mask[:, :, u:] *= frame_mask u2 = (u1 + self.encoder_dim[i]) // 2
feature_mask[:, :, u1:u2] *= frame_mask[..., 0:1]
feature_mask[:, :, u2:] *= frame_mask[..., 1:2]
feature_masks.append(feature_mask) feature_masks.append(feature_mask)
return feature_masks return feature_masks