本文的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 = size
self.tt = torch.cuda
self.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_step
self.nextYs[0][0] = sos
# Has EOS topped the beam yet.
self._eos = eos
self.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 batch
def 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_size
if 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] 刚开始只有第一个beam
flatBeamLk = 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] 记录了每个时间节点的结果来自于第几个beam
self.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设置为True
self.eosTop = True
def done(self):
# 当eosTop为True且已经结束的beam数大于等于beam_size的时候就结束。
return self.eosTop and len(self.finished) >=self.size
def 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是指该结果来自于第几个beam
hyp = []
for j in range(len(self.prevKs[:timestep]) - 1, -1, -1):
# prevKs: [time_step, beam_size] 记录了每个时间节点的结果来自于第几个beam
hyp.append(self.nextYs[j+1][k])
# nextYs: [time_step, beam_size] 记录了每个beam的每一步选择，将该id加入到hyp中
k = self.prevKs[j][k]
# k为结果来自于第几个beam
hyps.append(hyp[::-1])# hyp反过来加入到hyps中
# 最后得到的hyps：[beam_size, ~]列表，~即长度不一，是每一个beam的预测结果，按分数大小排列
return hyps
def buildTargetTokens(self, preds):
# preds即为getHyp产生的hyps，记录了每个beam产生的结果，按分数大小排列
# 这个函数的目的是截断eos之后的结果
sentence=[]
for pred in preds:
tokens = []
for tok in pred:
if tok==self._eos:
break
tokens.append(tok)
sentence.append(tokens)
return sentence