PyTorch torch.max over multiple dimensions

Question

Have tensor like :x.shape = [3, 2, 2].

import torch

x = torch.tensor([
    [[-0.3000, -0.2926],[-0.2705, -0.2632]],
    [[-0.1821, -0.1747],[-0.1526, -0.1453]],
    [[-0.0642, -0.0568],[-0.0347, -0.0274]]
])

我需要.max()接管第二和第三维。我期待一些这样的[-0.2632, -0.1453, -0.0274]输出。我尝试使用: x.max(dim=(1,2))，但这会导致错误。

Answer 1

现在，你可以做到这一点。该PR被合并（8月28日），现在可在夜间释放。

只需使用torch.amax()：

import torch

x = torch.tensor([
    [[-0.3000, -0.2926],[-0.2705, -0.2632]],
    [[-0.1821, -0.1747],[-0.1526, -0.1453]],
    [[-0.0642, -0.0568],[-0.0347, -0.0274]]
])

print(torch.amax(x, dim=(1, 2)))

# Output:
# >>> tensor([-0.2632, -0.1453, -0.0274])

原答案

截至今天（2020 年 4 月 11 日），在 PyTorch 中没有办法做.min()或.max()跨越多个维度。有一个关于它的公开问题，您可以关注它并查看它是否得到实施。您的情况的解决方法是：

import torch

x = torch.tensor([
    [[-0.3000, -0.2926],[-0.2705, -0.2632]],
    [[-0.1821, -0.1747],[-0.1526, -0.1453]],
    [[-0.0642, -0.0568],[-0.0347, -0.0274]]
])

print(x.view(x.size(0), -1).max(dim=-1))

# output:
# >>> values=tensor([-0.2632, -0.1453, -0.0274]),
# >>> indices=tensor([3, 3, 3]))

因此，如果您只需要值：x.view(x.size(0), -1).max(dim=-1).values.

如果x不是连续张量，.view()则将失败。在这种情况下，您应该.reshape()改用。

2020 年 8 月 26 日更新

此功能正在PR#43092 中实现，函数将被调用amin和amax。他们将只返回值。这可能很快就会被合并，所以当您阅读本文时，您可能能够在每晚构建时访问这些功能:) 玩得开心。

Answer 2

虽然Berriel 的解决方案解决了这个具体问题，但我认为添加一些解释可能有助于每个人对这里使用的技巧有所了解，以便它可以适用于 (m) 任何其他维度。

让我们从检查输入张量的形状开始x：

In [58]: x.shape   
Out[58]: torch.Size([3, 2, 2])

所以，我们有一个形状为 的 3D 张量(3, 2, 2)。现在，按照OP的问题，我们需要计算maximum的值在沿两个1张量^ST和2^次尺寸。在撰写本文时，torch.max()的dim论点仅支持int. 所以，我们不能使用元组。因此，我们将使用以下技巧，我将其称为，

平化＆最大招：因为我们需要计算max在两个1^日和2^次尺寸，我们将压平这两个层面，一个单一尺寸和离开0^个维度不变。这正是正在发生的事情：

In [61]: x.flatten().reshape(x.shape[0], -1).shape   
Out[61]: torch.Size([3, 4])   # 2*2 = 4

所以，现在我们已经将 3D 张量缩小为 2D 张量（即矩阵）。

In [62]: x.flatten().reshape(x.shape[0], -1) 
Out[62]:
tensor([[-0.3000, -0.2926, -0.2705, -0.2632],
        [-0.1821, -0.1747, -0.1526, -0.1453],
        [-0.0642, -0.0568, -0.0347, -0.0274]])

现在，我们可以简单地套用max在1^日的尺寸（即在这种情况下，第一个维度，也是最后一个维度），因为该维度扁平化规模居。

In [65]: x.flatten().reshape(x.shape[0], -1).max(dim=1)    # or: `dim = -1`
Out[65]: 
torch.return_types.max(
values=tensor([-0.2632, -0.1453, -0.0274]),
indices=tensor([3, 3, 3]))

由于矩阵中有 3 行，因此我们在合成张量中得到了 3 个值。

现在，在另一方面，如果你要计算max在0^次和1^日的尺寸，你会怎么做：

In [80]: x.flatten().reshape(-1, x.shape[-1]).shape 
Out[80]: torch.Size([6, 2])    # 3*2 = 6

In [79]: x.flatten().reshape(-1, x.shape[-1]) 
Out[79]: 
tensor([[-0.3000, -0.2926],
        [-0.2705, -0.2632],
        [-0.1821, -0.1747],
        [-0.1526, -0.1453],
        [-0.0642, -0.0568],
        [-0.0347, -0.0274]])

现在，我们可以简单地套用max在0^个维度，因为这是我们的扁平化的结果。（（另外，从我们的 ( 3, 2, 2) 的原始形状，在前 2 个维度上取最大值后，我们应该得到两个值作为结果。）

In [82]: x.flatten().reshape(-1, x.shape[-1]).max(dim=0) 
Out[82]: 
torch.return_types.max(
values=tensor([-0.0347, -0.0274]),
indices=tensor([5, 5]))

同样，您可以将此方法应用于多维和其他归约函数，例如min.

注意：我遵循基于 0 的维度 ( 0, 1, 2, 3, ...)的术语只是为了与 PyTorch 用法和代码保持一致。