gather_tree

  • paddle.fluid.layers.gather_tree(ids, parents)[源代码]

该OP在整个束搜索(Beam Search)结束后使用。在搜索结束后,可以获得每个时间步选择的的候选词id及其对应的在搜索树中的parent节点, idsparents 的形状布局均为

gather_tree - 图1 ,该OP从最后一个时间步回溯产生完整的id序列。

示例:

  1. 给定:
  2. ids = [[[2 2]
  3. [6 1]]
  4. [[3 9]
  5. [6 1]]
  6. [[0 1]
  7. [9 0]]]
  8. parents = [[[0 0]
  9. [1 1]]
  10. [[1 0]
  11. [1 0]]
  12. [[0 0]
  13. [0 1]]]
  14.  
  15. 结果:
  16. gather_tree(ids, parents)
  17. = [[[2 2]
  18. [1 6]]
  19. [[3 3]
  20. [6 1]]
  21. [[0 1]
  22. [9 0]]]
  • 参数:
    • ids (Variable) - 形状为 gather_tree - 图2 的三维Tensor,数据类型是int32或int64。包含了所有时间步选择的id。
    • parents (Variable) - 形状和数据类型均与 ids 相同的Tensor。包含了束搜索中每一时间步所选id对应的parent。

返回:和 ids 具有相同形状和数据类型的Tensor。包含了根据parent回溯而收集产生的完整id序列。

返回类型:Variable

代码示例

  1. import paddle.fluid as fluid
  2.  
  3. ids = fluid.data(name='ids',
  4. shape=[5, 2, 2],
  5. dtype='int64')
  6. parents = fluid.data(name='parents',
  7. shape=[5, 2, 2],
  8. dtype='int64')
  9. final_sequences = fluid.layers.gather_tree(ids, parents)