TensorFlow 返回張量維度上最大值的索引

tf.arg_max

arg_max (
     input ,
     dimension ,
     name = None
 )

返回在張量維度上具有最大值的索引.

請(qǐng)注意,在關(guān)聯(lián)的情況下,返回值的身份不能保證.

ARGS：

• input：張量.必須是下列類(lèi)型之一：float32,float64,int64,int32,uint8,uint16,int16,int8,complex64,complex128,qint8,quint8,qint32,half.
• dimension：張量.必須是以下類(lèi)型之一：int32,int64.當(dāng)類(lèi)型為 int32 時(shí),應(yīng)滿(mǎn)足：0 <= dimension <rank(input).描述輸入張量的哪個(gè)維度可以減少.對(duì)于向量,請(qǐng)使用dimension = 0.
• name：操作的名稱(chēng)(可選).

返回：

返回的張量類(lèi)型為 int64.

示例

argmax的源代碼

# pylint: disable=redefined-builtin
# TODO(aselle): deprecate arg_max
def argmax(input, axis=None, name=None, dimension=None):
  if dimension is not None:
    if axis is not None:
      raise ValueError("Cannot specify both 'axis' and 'dimension'")
    axis = dimension
  elif axis is None:
    axis = 0
  return gen_math_ops.arg_max(input, axis, name)

如您所見(jiàn)，argmax在內(nèi)部使用arg_max。另外，在代碼中，我建議使用argmax，因?yàn)閍rg_max可能很快就會(huì)被棄用。