From f57b68121a526baea43b2e93f4540d8a2995f633 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期一, 29 四月 2024 15:15:24 +0800
Subject: [PATCH] batch
---
funasr/models/sense_voice/model.py | 5 +++++
1 files changed, 5 insertions(+), 0 deletions(-)
diff --git a/funasr/models/sense_voice/model.py b/funasr/models/sense_voice/model.py
index ae20902..d06776b 100644
--- a/funasr/models/sense_voice/model.py
+++ b/funasr/models/sense_voice/model.py
@@ -310,6 +310,7 @@
speech_lengths = speech_lengths[:, 0]
batch_size, frames, _ = speech.shape
+ _, text_tokens = text.shape
if self.activation_checkpoint:
from torch.utils.checkpoint import checkpoint
@@ -331,6 +332,10 @@
stats["batch_size_x_frames"] = frames * batch_size
stats["batch_size_real_frames"] = speech_lengths.sum().item()
stats["padding_frames"] = stats["batch_size_x_frames"] - stats["batch_size_real_frames"]
+ stats["batch_size_x_tokens"] = text_tokens * batch_size
+ stats["batch_size_real_tokens"] = text_lengths.sum().item()
+ stats["padding_tokens"] = stats["batch_size_x_tokens"] - stats["batch_size_real_tokens"]
+ stats["batch_size_x_frames_plus_tokens"] = (text_tokens + frames) * batch_size
# force_gatherable: to-device and to-tensor if scalar for DataParallel
if self.length_normalized_loss:
--
Gitblit v1.9.1