speech_recognition/FunASR
1
0
Fork 0
mirror of https://github.com/modelscope/FunASR synced 2025-09-15 14:48:36 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

batch

Browse Source
This commit is contained in:
游雁 2024-04-29 15:15:24 +08:00
parent 2779602177
commit f57b68121a
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
funasr/models/sense_voice/model.py
View File

@ -310,6 +310,7 @@ class SenseVoiceRWKV(nn.Module):
speech_lengths = speech_lengths[:, 0] speech_lengths = speech_lengths[:, 0]
batch_size, frames, _ = speech.shape batch_size, frames, _ = speech.shape
_, text_tokens = text.shape
if self.activation_checkpoint: if self.activation_checkpoint:
from torch.utils.checkpoint import checkpoint from torch.utils.checkpoint import checkpoint
@ -331,6 +332,10 @@ class SenseVoiceRWKV(nn.Module):
stats["batch_size_x_frames"] = frames * batch_size stats["batch_size_x_frames"] = frames * batch_size
stats["batch_size_real_frames"] = speech_lengths.sum().item() stats["batch_size_real_frames"] = speech_lengths.sum().item()
stats["padding_frames"] = stats["batch_size_x_frames"] - stats["batch_size_real_frames"] 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 # force_gatherable: to-device and to-tensor if scalar for DataParallel
if self.length_normalized_loss: if self.length_normalized_loss: