Merge pull request #1809 from liugz18/main

Shi Xian

2024-06-17 ba54c2f88f8037f067d236febecb1e333b5e2543

 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