torch.argmax()

1. 官方文档：

torch.argmax(input, dim, keepdim=False) → LongTensor

Returns the indices of the maximum values of a tensor across a dimension.

Parameters

• input (Tensor) – the input tensor.
• dim (int) – the dimension to reduce. If None, the argmax of the flattened input is returned.
• keepdim (bool) – whether the output tensor has dim retained or not. Ignored if dim=None.

---

2.例子

x = torch.randn(2,4,5)

print(x)
> tensor([[[-0.2377, -1.0894,  1.7376, -1.7329,  0.4314],
         [-1.4899, -0.7765, -0.8258, -0.5007,  0.3901],
         [ 0.5550,  1.1423, -1.3531, -0.7470,  0.5840],
         [-0.7997,  0.0681,  0.7373, -0.0090, -0.9902]],

        [[-0.6789, -1.6443, -0.3977, -0.0783,  1.6161],
         [-1.5234, -1.1999, -0.9181,  0.8778,  0.5733],
         [ 1.1922, -0.4699, -0.2262,  0.9180,  0.0642],
         [-0.6872,  0.3727,  1.3596, -0.8487, -0.8330]]])

print(x.argmax(dim=0))
> tensor([[0, 0, 0, 1, 1],
        [0, 0, 0, 1, 1],
        [1, 0, 1, 1, 0],
        [1, 1, 1, 0, 1]])


print(x.argmax(dim=2))
> tensor([[2, 4, 1, 2],
        [4, 3, 0, 2]])

从上面两个例子可知，torch.argmax(dim)会返回dim维度上张量最大值的索引。

给 $x$ 一个具体的含义，$x=$[batch size, sentence len, vocal list dim]，vocal list 向量的每一个值，表示它是某个“单词”的概率。那选取网络预测的结果，就需要找到 [1,vocal list] 中的最大值的index，这个维度在代码中表示为 dim=2 。