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维度所有值进行取最大值。