torch.tensor.max 函数解析

pytorch 中的函数 torch.tensor.max(dim, keepdim = False)用于求某一维的最大值

返回值索引0为其数值，索引1为其最大值的索引

示例代码：

import torch

if __name__ = "__main__":
    test_tensor = torch.arange(6)
    test_tensor = test_tensor.view(3, 2)
    print('原始张量：', test_tensor)
    max_tensor = test_tensor.max(1, keepdim=True)
    print('max后0索引：', max_tensor[0])
    print('max后1索引：', max_tensor[1])
    print('max后全部返回值：', max_tensor)

输出结果：

原始张量： tensor(
       [[0, 1],
        [2, 3],
        [4, 5]])

max后0索引： tensor(
       [[1],
        [3],
        [5]])

max后1索引： tensor(
       [[1],
        [1],
        [1]])

max后全部返回值： torch.return_types.max(
values=tensor(
       [[1],
        [3],
        [5]]),
indices=tensor(
       [[1],
        [1],
        [1]]))

可见对一个3×2的矩阵的维度1进行max操作后，返回的value和indices都是3×1矩阵，计算过程如下：

value[0,0]=max(test_tensor[0,:])
value[1,0]=max(test_tensor[1,:])
value[2,0]=max(test_tensor[2,:])

可见对第dim维度进行max操作，就是把其它维度固定，对dim维度所有值进行取最大值。