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

Fix ONNX export of the latest streaming zipformer model. (#1148)

Browse Source
This commit is contained in:
Fangjun Kuang 2023-06-27 14:35:59 +08:00 committed by GitHub
parent 219bba1310
commit 968ebd236b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
egs/librispeech/ASR/zipformer/export-onnx-streaming.py
View File

@ -86,7 +86,7 @@ from icefall.checkpoint import (
find_checkpoints, find_checkpoints,
load_checkpoint, load_checkpoint,
) )
from icefall.utils import make_pad_mask, str2bool from icefall.utils import str2bool
def get_parser(): def get_parser():
@ -218,7 +218,7 @@ class OnnxEncoder(nn.Module):
) )
assert x.size(1) == self.chunk_size, (x.size(1), self.chunk_size) assert x.size(1) == self.chunk_size, (x.size(1), self.chunk_size)
src_key_padding_mask = make_pad_mask(x_lens) src_key_padding_mask = torch.zeros(N, self.chunk_size, dtype=torch.bool)
# processed_mask is used to mask out initial states # processed_mask is used to mask out initial states
processed_mask = torch.arange(left_context_len, device=x.device).expand( processed_mask = torch.arange(left_context_len, device=x.device).expand(
@ -272,6 +272,7 @@ class OnnxEncoder(nn.Module):
states = self.encoder.get_init_states(batch_size, device) states = self.encoder.get_init_states(batch_size, device)
embed_states = self.encoder_embed.get_init_states(batch_size, device) embed_states = self.encoder_embed.get_init_states(batch_size, device)
states.append(embed_states) states.append(embed_states)
processed_lens = torch.zeros(batch_size, dtype=torch.int64, device=device) processed_lens = torch.zeros(batch_size, dtype=torch.int64, device=device)