beam_search

paddle.fluid.layers.beam_search ( pre_ids, pre_scores, ids, scores, beam_size, end_id, level\=0, is_accumulated\=True, name\=None, return_parent_idx\=False ) [源代码]

束搜索(Beam search)是在机器翻译等生成任务中选择候选词的一种经典算法

更多细节参考 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完成。

束搜索的完整用法请参阅以下示例：

fluid/tests/book/test_machine_translation.py

参数:

  • pre_ids （Variable） - LoD level为2的LodTensor，表示前一时间步选择的候选id，是前一时间步 beam_search 的输出。第一步时，其形状应为为

    beam_search - 图1

    ， lod应为

    beam_search - 图2

    。数据类型为int64。

  • pre_scores （Variable） - 维度和LoD均与 pre_ids 相同的LodTensor，表示前一时间步所选id对应的累积得分，是前一时间步 beam_search 的输出。数据类型为float32。

  • ids （None|Variable） - 包含候选id的LodTensor。LoD应与 pre_ids 相同，形状为

    beam_search - 图3

    ，其中第一维大小与 pre_ids 相同且``batch_size`` 会随样本到达结束而自动减小， K 应该大于 beam_size 。数据类型为int64。可为空，为空时使用 scores 上的索引作为id。

  • scores （Variable） - 表示 ids 对应的累积分数的LodTensor变量, 维度和LoD均与 ids 相同。

  • beam_size （int） - 指明束搜索中的束宽度。

  • end_id （int） - 指明标识序列结束的id。

  • level （int，可选） - 可忽略，当前不能更改 。知道LoD level为2即可，两层lod的意义如下： 第一级表示每个源句（样本）包含的beam大小，若满足结束条件（达到 beam_size 个结束）则变为0；第二级是表示每个beam被选择的次数。

  • is_accumulated （bool，可选） - 指明输入分数 scores 是否为累积分数，默认为True。

  • name (str，可选) – 具体用法请参见 Name ，一般无需设置，默认值为None。

  • return_parent_idx （bool，可选） - 指明是否返回一个额外的Tensor，该Tensor保存了选择的id的父节点（beam）在 pre_id 中索引，可用于通过gather OP更新其他Tensor的内容。默认为False。

返回：Variable的二元组或三元组。二元组中包含了当前时间步选择的id和对应的累积得分两个LodTensor，形状相同且均为 [batch_size×beam_size，1][batch_size×beam_size，1] ，LoD相同且level均为2，数据类型分别为int64和float32；若 return_parent_idx 为True时为三元组，多返回一个保存了父节点在 pre_id 中索引的Tensor，形状为 [batch_size×beam_size][batch_size×beam_size] ，数据类型为int64。

返回类型：tuple

代码示例

  1. import paddle.fluid as fluid
  3. # 假设 `probs` 包含计算神经元所得的预测结果
  4. # `pre_ids` 和 `pre_scores` 为beam_search之前时间步的输出
  5. beam_size = 4
  6. end_id = 1
  7. pre_ids = fluid.layers.data(
  8.     name='pre_id', shape=[1], lod_level=2, dtype='int64')
  9. pre_scores = fluid.layers.data(
  10.     name='pre_scores', shape=[1], lod_level=2, dtype='float32')
  11. probs = fluid.layers.data(
  12.     name='probs', shape=[10000], dtype='float32')
  13. topk_scores, topk_indices = fluid.layers.topk(probs, k=beam_size)
  14. accu_scores = fluid.layers.elementwise_add(
  15.                                       x=fluid.layers.log(x=topk_scores),
  16.                                       y=fluid.layers.reshape(
  17.                                           pre_scores, shape=[-1]),
  18.                                       axis=0)
  19. selected_ids, selected_scores = fluid.layers.beam_search(
  20.                                       pre_ids=pre_ids,
  21.                                       pre_scores=pre_scores,
  22.                                       ids=topk_indices,
  23.                                       scores=accu_scores,
  24.                                       beam_size=beam_size,
  25.                                       end_id=end_id)