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)
全部评论

相关推荐

点赞 收藏 评论
分享
牛客网
牛客企业服务