From ba54c2f88f8037f067d236febecb1e333b5e2543 Mon Sep 17 00:00:00 2001
From: Shi Xian <40013335+R1ckShi@users.noreply.github.com>
Date: 星期一, 17 六月 2024 11:11:25 +0800
Subject: [PATCH] Merge pull request #1809 from liugz18/main
---
funasr/models/paraformer/cif_predictor.py | 33 +++++++++++++++++++++++----------
1 files changed, 23 insertions(+), 10 deletions(-)
diff --git a/funasr/models/paraformer/cif_predictor.py b/funasr/models/paraformer/cif_predictor.py
index a6bfe65..05e283a 100644
--- a/funasr/models/paraformer/cif_predictor.py
+++ b/funasr/models/paraformer/cif_predictor.py
@@ -504,7 +504,7 @@
frames = torch.zeros(batch_size, len_time, hidden_size, dtype=dtype, device=device)
fires = torch.zeros(batch_size, len_time, dtype=dtype, device=device)
- prefix_sum = torch.cumsum(alphas, dim=1)
+ prefix_sum = torch.cumsum(alphas, dim=1, dtype=torch.float64).to(torch.float32) # cumsum precision degradation cause wrong result in extreme
prefix_sum_floor = torch.floor(prefix_sum)
dislocation_prefix_sum = torch.roll(prefix_sum, 1, dims=1)
dislocation_prefix_sum_floor = torch.floor(dislocation_prefix_sum)
@@ -539,7 +539,8 @@
frames = frames - shift_frames + shift_remain_frames - remain_frames
- max_label_len = batch_len.max()
+ max_label_len = alphas.sum(dim=-1)
+ max_label_len = torch.floor(max_label_len).max().to(dtype=torch.int64)
frame_fires = torch.zeros(
batch_size, max_label_len, hidden_size, dtype=dtype, device=device
@@ -661,17 +662,16 @@
return torch.stack(list_ls, 0), fires
-def cif_v1(hidden, alphas, threshold):
+def cif_wo_hidden_v1(alphas, threshold, return_fire_idxs=False):
+ batch_size, len_time = alphas.size()
+ device = alphas.device
+ dtype = alphas.dtype
- device = hidden.device
- dtype = hidden.dtype
- batch_size, len_time, hidden_size = hidden.size()
threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device)
- frames = torch.zeros(batch_size, len_time, hidden_size, dtype=dtype, device=device)
fires = torch.zeros(batch_size, len_time, dtype=dtype, device=device)
- prefix_sum = torch.cumsum(alphas, dim=1)
+ prefix_sum = torch.cumsum(alphas, dim=1, dtype=torch.float64).to(torch.float32) # cumsum precision degradation cause wrong result in extreme
prefix_sum_floor = torch.floor(prefix_sum)
dislocation_prefix_sum = torch.roll(prefix_sum, 1, dims=1)
dislocation_prefix_sum_floor = torch.floor(dislocation_prefix_sum)
@@ -682,7 +682,19 @@
fire_idxs = dislocation_diff > 0
fires[fire_idxs] = 1
fires = fires + prefix_sum - prefix_sum_floor
+ if return_fire_idxs:
+ return fires, fire_idxs
+ return fires
+
+def cif_v1(hidden, alphas, threshold):
+ fires, fire_idxs = cif_wo_hidden_v1(alphas, threshold, return_fire_idxs=True)
+
+ device = hidden.device
+ dtype = hidden.dtype
+ batch_size, len_time, hidden_size = hidden.size()
+ frames = torch.zeros(batch_size, len_time, hidden_size,
+ dtype=dtype, device=device)
prefix_sum_hidden = torch.cumsum(
alphas.unsqueeze(-1).tile((1, 1, hidden_size)) * hidden, dim=1
)
@@ -698,7 +710,8 @@
remains = fires - torch.floor(fires)
remain_frames = (
- remains[fire_idxs].unsqueeze(-1).tile((1, hidden_size)) * hidden[fire_idxs]
+ remains[fire_idxs].unsqueeze(-1).tile((1,
+ hidden_size)) * hidden[fire_idxs]
)
shift_remain_frames = torch.roll(remain_frames, 1, dims=0)
@@ -706,7 +719,7 @@
frames = frames - shift_frames + shift_remain_frames - remain_frames
- max_label_len = batch_len.max()
+ max_label_len = torch.round(alphas.sum(-1)).int().max() # torch.round to calculate the max length
frame_fires = torch.zeros(
batch_size, max_label_len, hidden_size, dtype=dtype, device=device
--
Gitblit v1.9.1