Beam Search源码理解

本文的beam search源码来自：CodeBERT/model.py at master · microsoft/CodeBERT (github.com)https://github.com/microsoft/CodeBERT/blob/master/CodeBERT/code2nl/model.py

理解过程中加入了注释：

class Beam(object):def __init__(self, size,sos,eos):self.size = sizeself.tt = torch.cudaself.scores = self.tt.FloatTensor(size).zero_()# 大小为[beam_size]，记录当前每个beam的分数总和self.prevKs = []# 记录每一步选取的是第几个beam，便于最后回溯生成结果self.nextYs = [self.tt.LongTensor(size).fill_(0)]# nextYs: [seq_len=1, beam_size]，随着预测过程seq_len逐渐增加，表示每一步的输出结果# seq_len即为time_stepself.nextYs[0][0] = sos# Has EOS topped the beam yet.self._eos = eosself.eosTop = False# Time and k pair for finished.self.finished = []def getCurrentState(self):batch = self.tt.LongTensor(self.nextYs[-1]).view(-1, 1)# batch: [beam_size, seq_len]，用于加入到下一次模型的输入中。return batchdef getCurrentOrigin(self):"Get the backpointers for the current timestep."return self.prevKs[-1]def advance(self, wordLk):'''更新beam中的信息wordLk: [beam_size, vocab_size]，上一个时间节点每个beam的模型预测结果，需要用LogSoftMax进行归一化'''numWords = wordLk.size(1)# numWords: vocab_sizeif len(self.prevKs) > 0:beamLk = wordLk + self.scores.unsqueeze(1).expand_as(wordLk)# scores: [beam_size]# wordLk是当前的分数，scores是之前的分数，加起来得到beamLk: [beam_size, vocab_size]for i in range(self.nextYs[-1].size(0)):if self.nextYs[-1][i] == self._eos:beamLk[i] = -1e20# 把第i个beam的概率全部设置为负无穷else:beamLk = wordLk[0]# beamLk: [vocab_size] 刚开始只有第一个beamflatBeamLk = beamLk.view(-1) # beamlLk展开bestScores, bestScoresId = flatBeamLk.topk(self.size, 0, True, True) # topk个最好分数self.scores = bestScores# scores: [beam_size]prevK = bestScoresId // numWords# prevK: [beam_size]self.prevKs.append(prevK)# prevKs: [time_step, beam_size] 记录了每个时间节点的结果来自于第几个beamself.nextYs.append((bestScoresId - prevK * numWords))# nextYs: [seq_len, beam_size] 记录了每个事件节点选取的id, seq_len即time_step# 对nextYs的最后一个时间节点进行遍历，检查是否出现了结束符for i in range(self.nextYs[-1].size(0)):if self.nextYs[-1][i] == self._eos:s = self.scores[i]self.finished.append((s, len(self.nextYs) - 1, i))# i 表示第几个beam# 若出现结束符，将(总分数，句子长度，beam的id)三元组加入到finished列表中# finished列表中存的是已经结束的beam的信息# End condition is when top-of-beam is EOS and no global score.if self.nextYs[-1][0] == self._eos:# 当nextYs中最后一个时间点的第一个id为结束符时，将eosTop设置为Trueself.eosTop = Truedef done(self):# 当eosTop为True且已经结束的beam数大于等于beam_size的时候就结束。return self.eosTop and len(self.finished) >=self.sizedef getFinal(self):if len(self.finished) == 0:# 这里的情况就是所有beam的句子长度都达到了max_length但没有任何一个产生了结束符self.finished.append((self.scores[0], len(self.nextYs) - 1, 0))# 这种情况下就手动将第0个beam设置为已经结束self.finished.sort(key=lambda a: -a[0])# 将finished按beam的分数由大到小排序if len(self.finished) != self.size:# 将没有结束的句子也按(分数，长度，beam_id)三元组的形势加入到finished中unfinished=[]for i in range(self.nextYs[-1].size(0)):if self.nextYs[-1][i] != self._eos:s = self.scores[i]unfinished.append((s, len(self.nextYs) - 1, i)) unfinished.sort(key=lambda a: -a[0])self.finished+=unfinished[:self.size-len(self.finished)]# 已经结束的beam排在未结束的句子前面return self.finished[:self.size]def getHyp(self, beam_res):"""回溯，生成结果"""# beam_res 传入的就是finished列表，由get_final得到hyps=[]for _,timestep, k in beam_res:# k是指该结果来自于第几个beamhyp = []for j in range(len(self.prevKs[:timestep]) - 1, -1, -1):# prevKs: [time_step, beam_size] 记录了每个时间节点的结果来自于第几个beamhyp.append(self.nextYs[j+1][k])# nextYs: [time_step, beam_size] 记录了每个beam的每一步选择，将该id加入到hyp中k = self.prevKs[j][k]# k为结果来自于第几个beamhyps.append(hyp[::-1])# hyp反过来加入到hyps中# 最后得到的hyps：[beam_size, ~]列表，~即长度不一，是每一个beam的预测结果，按分数大小排列return hypsdef buildTargetTokens(self, preds):# preds即为getHyp产生的hyps，记录了每个beam产生的结果，按分数大小排列# 这个函数的目的是截断eos之后的结果sentence=[]for pred in preds:tokens = []for tok in pred:if tok==self._eos:breaktokens.append(tok)sentence.append(tokens)return sentence

Beam Search源码理解相关推荐

压缩跟踪Compressive Tracking源码理解
压缩跟踪Compressive Tracking源码理解 zouxy09@qq.com http://blog.csdn.net/zouxy09 在前面一个介绍<Real-Time Compre ...
从hotspot底层对象结构理解锁膨胀升级过程||深入jdk源码理解longadder的分段cas优化机制——分段CAS优化
深入jdk源码理解longadder的分段cas优化机制 longadder
faster rcnn源码理解（二）之AnchorTargetLayer（网络中的rpn_data）
转载自:faster rcnn源码理解(二)之AnchorTargetLayer(网络中的rpn_data) - 野孩子的专栏 - 博客频道 - CSDN.NET http://blog.csdn.n ...
faster rcnn的源码理解（一）SmoothL1LossLayer论文与代码的结合理解
转载自:faster rcnn的源码理解(一)SmoothL1LossLayer论文与代码的结合理解 - 野孩子的专栏 - 博客频道 - CSDN.NET http://blog.csdn.net/u ...
TLD（Tracking-Learning-Detection）学习与源码理解之（六）
TLD(Tracking-Learning-Detection)学习与源码理解之(六) zouxy09@qq.com http://blog.csdn.net/zouxy09 下面是自己在看论文和这些 ...
TLD（Tracking-Learning-Detection）学习与源码理解之（五）
TLD(Tracking-Learning-Detection)学习与源码理解之(五) zouxy09@qq.com http://blog.csdn.net/zouxy09 下面是自己在看论文和 ...
TLD（Tracking-Learning-Detection）学习与源码理解之（四）
TLD(Tracking-Learning-Detection)学习与源码理解之(四) zouxy09@qq.com http://blog.csdn.net/zouxy09 下面是自己在看论文和这些 ...
TLD（Tracking-Learning-Detection）学习与源码理解之（三）
TLD(Tracking-Learning-Detection)学习与源码理解之(三) zouxy09@qq.com http://blog.csdn.net/zouxy09 下面是自己在看论文和这些 ...
TLD（Tracking-Learning-Detection）学习与源码理解之（二）
TLD(Tracking-Learning-Detection)学习与源码理解之(二) zouxy09@qq.com http://blog.csdn.net/zouxy09 OpenTLD下载与编译 ...

Beam Search源码理解

Beam Search源码理解相关推荐

最新文章

热门文章