fix bug about memory mask when memory_size==0

This commit is contained in:
yaozengwei 2022-07-04 17:05:47 +08:00
parent 9c37c16326
commit 10662c5c38
3 changed files with 32 additions and 17 deletions

View File

@ -9,8 +9,7 @@ per-file-ignores =
egs/*/ASR/pruned_transducer_stateless*/*.py: E501,
egs/*/ASR/*/optim.py: E501,
egs/*/ASR/*/scaling.py: E501,
egs/librispeech/ASR/conv_emformer_transducer_stateless/*.py: E501, E203
egs/librispeech/ASR/conv_emformer_transducer_stateless2/*.py: E501, E203
egs/librispeech/ASR/conv_emformer_transducer_stateless*/*.py: E501, E203
# invalid escape sequence (cause by tex formular), W605
icefall/utils.py: E501, W605

View File

@ -1579,13 +1579,19 @@ class EmformerEncoder(nn.Module):
# calcualte padding mask to mask out initial zero caches
chunk_mask = make_pad_mask(output_lengths).to(x.device)
memory_mask = (
torch.div(
num_processed_frames, self.chunk_length, rounding_mode="floor"
).view(x.size(1), 1)
<= torch.arange(self.memory_size, device=x.device).expand(
x.size(1), self.memory_size
)
).flip(1)
(
torch.div(
num_processed_frames,
self.chunk_length,
rounding_mode="floor",
).view(x.size(1), 1)
<= torch.arange(self.memory_size, device=x.device).expand(
x.size(1), self.memory_size
)
).flip(1)
if self.use_memory
else torch.empty(0).to(dtype=torch.bool, device=x.device)
)
left_context_mask = (
num_processed_frames.view(x.size(1), 1)
<= torch.arange(self.left_context_length, device=x.device).expand(

View File

@ -1388,7 +1388,11 @@ class EmformerEncoder(nn.Module):
output_lengths = torch.clamp(lengths - self.right_context_length, min=0)
attention_mask = self._gen_attention_mask(utterance)
M = right_context.size(0) // self.right_context_length - 1
M = (
right_context.size(0) // self.right_context_length - 1
if self.use_memory
else 0
)
padding_mask = make_pad_mask(M + right_context.size(0) + output_lengths)
output = utterance
@ -1480,13 +1484,19 @@ class EmformerEncoder(nn.Module):
# calcualte padding mask to mask out initial zero caches
chunk_mask = make_pad_mask(output_lengths).to(x.device)
memory_mask = (
torch.div(
num_processed_frames, self.chunk_length, rounding_mode="floor"
).view(x.size(1), 1)
<= torch.arange(self.memory_size, device=x.device).expand(
x.size(1), self.memory_size
)
).flip(1)
(
torch.div(
num_processed_frames,
self.chunk_length,
rounding_mode="floor",
).view(x.size(1), 1)
<= torch.arange(self.memory_size, device=x.device).expand(
x.size(1), self.memory_size
)
).flip(1)
if self.use_memory
else torch.empty(0).to(dtype=torch.bool, device=x.device)
)
left_context_mask = (
num_processed_frames.view(x.size(1), 1)
<= torch.arange(self.left_context_length, device=x.device).expand(