argmax

paddle.fluid.layers. argmax ( x, axis=0 ) [源代码]

argmax

该OP沿 axis 计算输入 x 的最大元素的索引。

参数:

  • x (Variable) - 输入的多维 Tensor ,支持的数据类型:float32、float64、int8、int16、int32、int64。

  • axis (int,可选) - 指定对输入Tensor进行运算的轴, axis 的有效范围是[-R, R),R是输入 x 的Rank, axis 为负时与 axis +R 等价。默认值为0。

返回: Tensor ,数据类型int64

返回类型:Variable

代码示例

  1. import paddle.fluid as fluid
  2. import numpy as np
  3. in1 = np.array([[[5,8,9,5],
  4. [0,0,1,7],
  5. [6,9,2,4]],
  6. [[5,2,4,2],
  7. [4,7,7,9],
  8. [1,7,0,6]]])
  9. with fluid.dygraph.guard():
  10. x = fluid.dygraph.to_variable(in1)
  11. out1 = fluid.layers.argmax(x=x, axis=-1)
  12. out2 = fluid.layers.argmax(x=x, axis=0)
  13. out3 = fluid.layers.argmax(x=x, axis=1)
  14. out4 = fluid.layers.argmax(x=x, axis=2)
  15. print(out1.numpy())
  16. # [[2 3 1]
  17. # [0 3 1]]
  18. print(out2.numpy())
  19. # [[0 0 0 0]
  20. # [1 1 1 1]
  21. # [0 0 0 1]]
  22. print(out3.numpy())
  23. # [[2 2 0 1]
  24. # [0 1 1 1]]
  25. print(out4.numpy())
  26. # [[2 3 1]
  27. # [0 3 1]]