torch.max()函数

1️⃣ torch.max 的核心规则

torch.max(x, dim = k)
👉 固定除第 k 层以外的所有位置，只在第 k 层上比较大小
👉 第 k 层会被“压掉”

2️⃣ 示例张量

x = [
        [
            [ 1,  9,  3],
            [ 4,  2,  8],
        ],
        [
            [ 5,  0,  6],
            [ 7, 10,  2],
        ]
    ]

这是一个3维张量，形状为(2, 2, 3)，我们将对其进行torch.max操作。

3️⃣ torch.max(x, dim = 0)

dim=0代表在第0层上比较大小，即比较两个2x3的矩阵对应位置的元素，取较大值。结果是一个2x3的矩阵，形状为(2, 3)。

[0][0][0] = 1
[1][0][0] = 5
1 vs 5, 取较大值5

[0][0][1] = 9
[1][0][1] = 0
9 vs 0, 取较大值9

[0][0][2] = 3
[1][0][2] = 6
3 vs 6, 取较大值6

[0][1][0] = 4
[1][1][0] = 7
4 vs 7, 取较大值7

[0][1][1] = 2
[1][1][1] = 10
2 vs 10, 取较大值10

[0][1][2] = 8
[1][1][2] = 2
8 vs 2, 取较大值8

结果：

1 2	[[5, 9, 6], [7, 10, 8]]

4️⃣ torch.max(x, dim = 1)

dim=1代表在第1层上比较大小，即比较两个2x3的矩阵的行，取较大值。结果是一个2x3的矩阵，形状为(2, 3)。

[0][0][0] = 1
[0][1][0] = 4
1 vs 4, 取较大值4

[0][0][1] = 9
[0][1][1] = 2
9 vs 2, 取较大值9

[0][0][2] = 3
[0][1][2] = 8
3 vs 8, 取较大值8

[1][0][0] = 5
[1][1][0] = 7
5 vs 7, 取较大值7

[1][0][1] = 0
[1][1][1] = 10
0 vs 10, 取较大值10

[1][0][2] = 6
[1][1][2] = 2
6 vs 2, 取较大值6

结果：

1 2	[[4, 9, 8], [7, 10, 6]]

5️⃣ torch.max(x, dim = 2)

dim=2代表在第2层上比较大小，即比较两个2x3的矩阵的列，取较大值。结果是一个2x2的矩阵，形状为(2, 2)。

[0][0][0] = 1
[0][0][1] = 9
[0][0][2] = 3
1 vs 9 vs 3, 取较大值9

[0][1][0] = 4
[0][1][1] = 2
[0][1][2] = 8
4 vs 2 vs 8, 取较大值8

[1][0][0] = 5
[1][0][1] = 0
[1][0][2] = 6
5 vs 0 vs 6, 取较大值6

[1][1][0] = 7
[1][1][1] = 10
[1][1][2] = 2
7 vs 10 vs 2, 取较大值10

结果：

1 2	[[9, 8], [6, 10]]