torch.argmax¶
- torch.argmax(input) LongTensor ¶
返回
input
张量中所有元素的最大值的索引。这是
torch.max()
返回的第二个值。有关此方法的确切语义,请参阅其文档。注意
如果存在多个最大值,则返回第一个最大值的索引。
- 参数
input (Tensor) – 输入张量。
示例
>>> a = torch.randn(4, 4) >>> a tensor([[ 1.3398, 0.2663, -0.2686, 0.2450], [-0.7401, -0.8805, -0.3402, -1.1936], [ 0.4907, -1.3948, -1.0691, -0.3132], [-1.6092, 0.5419, -0.2993, 0.3195]]) >>> torch.argmax(a) tensor(0)
- torch.argmax(input, dim, keepdim=False) LongTensor
返回跨维度张量中最大值的索引。
这是
torch.max()
返回的第二个值。有关此方法的确切语义,请参阅其文档。- 参数
示例
>>> a = torch.randn(4, 4) >>> a tensor([[ 1.3398, 0.2663, -0.2686, 0.2450], [-0.7401, -0.8805, -0.3402, -1.1936], [ 0.4907, -1.3948, -1.0691, -0.3132], [-1.6092, 0.5419, -0.2993, 0.3195]]) >>> torch.argmax(a, dim=1) tensor([ 0, 2, 0, 1])