CTC束搜索解码原理和Pytorch实现（CTC Prefix BeamSearch Decode)

原创已于 2022-10-13 10:39:43 修改 · 2.1k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#python #算法 #人工智能

于 2022-10-12 17:02:51 首次发布

本文介绍CTC解码算法的工作原理及实现细节，包括如何通过原生序列生成标签序列，以及在推断过程中如何利用beam search进行高效搜索。文章详细解释了在不同情况下原生序列的增长策略，并给出了具体的实现代码。

CTC解码在推断时，同一个标签序列对应的原生序列的结尾会有两种情况：1.以字符结尾；2.以blank结尾。不同的结尾往下增长时的缩放策略不同，比如以字符结尾：*a遇到a会缩放为*a；以blank（用“-”表示）结尾：*a-遇到a会被缩放为*aa。所以在增长过程的每一步，标签序列的概率都会使用两个变量存储，一个负责累加以字符结尾的原生序列概率，另一个负责累加以blank结尾的原生序列概率，两者相互独立，互无交集。增长后，再将这两个概率相加(log_sum_exp)表示这一个标签序列的总概率。然后取top beam_size后再往下增长。

序列增长时会有四种情况:

原生序列结尾任意，当前值为blank，标签序列不变，更新以blank结尾的概率；
原生序列结尾为blank，当前值为相同字符（指与目前标签序列的最后一个字符相同），标签序列更新，更新非blank概率；
原生序列结尾为字符，当前值为相同字符，标签序列不变，更新非blank概率；
原生序列结尾任意，当前值为不同字符，标签序列更新，更新非blank概率。

注：

1.原生序列是指未缩放的序列，如aa-bbc-，aabbcc 对应的标签序列都为abc。

2.这里的概率指得都是对数概率：lp=log(softmax(logits))。所以原生序列增长时，其概率lp用“+”更新，相当于概率积后取log。而原生序列和标签序列是多对一关系，同一个标签序列的概率用其对应的多个原生序列概率的log_sum_exp表示(log(exp(lp1)+exp(lp2),...exp(lpk))，相当于概率和后再规范为对数概率表示。

import math

def log_sum_exp(lps):
    _inf = -float('inf')
    if all(lp == _inf for lp in lps):return _inf
    mlp = max(lps)
    return mlp + math.log(sum(math.exp(lp - mlp) for lp in lps))


def beam_search_ctc(probs,bms=10,blank=0):
    '''
    probs: 概率空间，shape为[sequence_len,vocab_size]的torch tensor
    bms: beam_size
    blank: blank index
    '''
    _inf = -float("inf")
    seqs =[((idx.item(),),(lp.item(),_inf)) if idx.item()!=blank
           else (tuple(),(_inf,lp.item()))
           for lp,idx in zip(*probs[0].topk(bms))]

    for i in range(1,probs.size(0)):
        new_seqs = {}
        for seq,(lps,blps) in seqs:   
            last = seq[-1] if len(seq) > 0 else None
            for lp, idx in zip(*probs[i].topk(bms)):
                lp=lp.item()
                idx=idx.item()  
                if idx == blank :
                    nlps,nblps= new_seqs.get(seq,(_inf,_inf))
                    new_seqs[seq]=(nlps,log_sum_exp([nblps,lps+lp,blps+lp]))
                elif idx ==last:
                    #aa
                    nlps,nblps= new_seqs.get(seq,(_inf,_inf))
                    new_seqs[seq]=(log_sum_exp([nlps,lps+lp]),nblps)
                    #a-a
                    new_seq = seq + (idx,)
                    nlps,nblps= new_seqs.get(new_seq,(_inf,_inf))
                    new_seqs[new_seq]=(log_sum_exp([nlps,blps+lp]),nblps)
                else:
                    new_seq = seq + (idx,)
                    nlps,nblps= new_seqs.get(new_seq,(_inf,_inf))
                    new_seqs[new_seq]=(log_sum_exp([nlps,lps+lp,blps+lp]),nblps)
        new_seqs = sorted(
            new_seqs.items(),
            key=lambda x: log_sum_exp(list(x[1])),
            reverse=True)
        seqs = new_seqs[:bms]     
    return seqs