beam search
该OP仅支持LoDTensor,在计算产生得分之后使用,完成单个时间步内的束搜索。具体而言,在计算部分产生 ids 和 scores 后,对于每个源句(样本)该OP从 ids 中根据其对应的 scores 选择当前时间步 top-K (K 是 beam_size)的候选词id。而 pre_id 和 pre_scores 是上一时间步 beam_search 的输出,加入输入用于特殊处理到达结束的翻译候选。
注意,如果 is_accumulated 为 True,传入的 scores 应该是累积分数。反之,scores 是单步得分,会在该OP内被转化为log值并累积到 pre_scores 作为最终得分。如需使用长度惩罚,应在计算累积分数前使用其他OP完成。
import paddle.fluid as fluid # 假设 `probs` 包含计算神经元所得的预测结果 # `pre_ids` 和 `pre_scores` 为beam_search之前时间步的输出 beam_size = 4 end_id = 1 pre_ids = fluid.layers.data( name='pre_id', shape=[1], lod_level=2, dtype='int64') pre_scores = fluid.layers.data( name='pre_scores', shape=[1], lod_level=2, dtype='float32') probs = fluid.layers.data( name='probs', shape=[10000], dtype='float32') topk_scores, topk_indices = fluid.layers.topk(probs, k=beam_size) accu_scores = fluid.layers.elementwise_add( x=fluid.layers.log(x=topk_scores), y=fluid.layers.reshape( pre_scores, shape=[-1]), axis=0) selected_ids, selected_scores = fluid.layers.beam_search( pre_ids=pre_ids, pre_scores=pre_scores, ids=topk_indices, scores=accu_scores, beam_size=beam_size, end_id=end_id)