mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-08-08 09:32:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove if-branch at downsample pad in zipformer for onnx-export compatibility (#965)

Browse Source
This commit is contained in:
Peng He 2023-03-24 14:30:43 +08:00 committed by GitHub
parent d74822d07b
commit f260a09ed4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -781,13 +781,12 @@ class AttentionDownsample(torch.nn.Module):
ds = self.downsample
d_seq_len = (seq_len + ds - 1) // ds
# Pad to an exact multiple of self.downsample
if seq_len != d_seq_len * ds:
# right-pad src, repeating the last element.
pad = d_seq_len * ds - seq_len
src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2])
src = torch.cat((src, src_extra), dim=0)
assert src.shape[0] == d_seq_len * ds, (src.shape[0], d_seq_len, ds)
# Pad to an exact multiple of self.downsample, could be 0 for onnx-export-compatibility
# right-pad src, repeating the last element.
pad = d_seq_len * ds - seq_len
src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2])
src = torch.cat((src, src_extra), dim=0)
assert src.shape[0] == d_seq_len * ds, (src.shape[0], d_seq_len, ds)
src = src.reshape(d_seq_len, ds, batch_size, in_channels)
scores = (src * self.query).sum(dim=-1, keepdim=True)