From 01b744f1279420e7f1017368dc6f120a61f04870 Mon Sep 17 00:00:00 2001 From: marcoyang Date: Sun, 7 Apr 2024 15:45:28 +0800 Subject: [PATCH] support onnx export with batch size 1; also works for batch processing, but the results might be affected by the padding --- egs/audioset/AT/zipformer/export-onnx.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/egs/audioset/AT/zipformer/export-onnx.py b/egs/audioset/AT/zipformer/export-onnx.py index 24bd431fc..af83c0e9c 100755 --- a/egs/audioset/AT/zipformer/export-onnx.py +++ b/egs/audioset/AT/zipformer/export-onnx.py @@ -202,10 +202,7 @@ class OnnxAudioTagger(nn.Module): logits = self.classifier(encoder_out) # (N, T, num_classes) # Note that this is slightly different from model.py for better # support of onnx - N = logits.shape[0] - for i in range(N): - logits[i, encoder_out_lens[i] :] = 0 - logits = logits.sum(dim=1) / encoder_out_lens.unsqueeze(-1) + logits = logits.mean(dim=1) return logits