如何有效地检索 Torch 张量中最大值的索引？

Posted 2023-02-14

【中文标题】如何有效地检索 Torch 张量中最大值的索引？

【英文标题】：How to efficiently retrieve the indices of maximum values in a Torch tensor?

【发布时间】：2019-04-12 05:49:54

【问题描述】：

假设有一个火炬张量，例如以下形状：

x = torch.rand(20, 1, 120, 120)

我现在想要的是获取每个 120x120 矩阵的最大值的索引。为了简化问题，我首先将x.squeeze() 使用形状[20, 120, 120]。然后我想得到火炬张量，它是一个形状为[20, 2]的索引列表。

我怎样才能快速做到这一点？

【问题讨论】：

为什么需要[20, 2] 矩阵。您是否想要每个 120 * 120 矩阵的行最大值和列最大值？

是的，或者换句话说：对于 20 个120 * 120 矩阵中的每一个，我想要具有最大值的单元格的[x, y] 坐标

如果你想知道***k元素的索引，使用torch.topk()。

这能回答你的问题吗？ Extracting the top-k value-indices from a 1-D Tensor

【参考方案1】：

torch.topk() 是您正在寻找的。从文档中，

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)

返回给定input张量的k最大元素 给定的维度。

如果没有给出dim，则选择输入的最后一个维度。

如果largest 是False，则返回k 个最小的元素。

返回一个 (values, indices) 的命名元组，其中索引是原始输入张量中元素的索引。

布尔选项sorted如果True，将确保返回的k个元素本身是排序的

【讨论】：

有用的功能知道，但它没有回答原来的问题。 OP 希望获得该矩阵中最大元素的 20 个 120x120 矩阵中的每一个的索引。也就是说，她想要 20 个 2D 坐标，每个矩阵一个。 topk 仅返回最大化维度中最大元素的索引。

请注意，topk 的文档在返回索引的含义方面存在混淆。它给人的印象是该函数为原始张量提供索引，而实际上它仅返回最大化维度的索引。请参阅试图澄清它的 pytorch 问题github.com/pytorch/pytorch/issues/50331#issue-782748956。

【参考方案2】：

如果我理解正确，您不需要值，而是需要索引。不幸的是，没有开箱即用的解决方案。有一个 argmax() 函数，但我看不出如何让它完全按照你的意愿去做。

所以这里有一个小变通方法，因为我们只是划分张量，所以效率应该还可以：

n = torch.tensor(4)
d = torch.tensor(4)
x = torch.rand(n, 1, d, d)
m = x.view(n, -1).argmax(1)
# since argmax() does only return the index of the flattened
# matrix block we have to calculate the indices by ourself 
# by using / and % (// would also work, but as we are dealing with
# type torch.long / works as well
indices = torch.cat(((m / d).view(-1, 1), (m % d).view(-1, 1)), dim=1)
print(x)
print(indices)

n 代表您的第一个维度，d 代表最后两个维度。我在这里取较小的数字来显示结果。但当然这也适用于n=20 和d=120：

n = torch.tensor(20)
d = torch.tensor(120)
x = torch.rand(n, 1, d, d)
m = x.view(n, -1).argmax(1)
indices = torch.cat(((m / d).view(-1, 1), (m % d).view(-1, 1)), dim=1)
#print(x)
print(indices)

这是n=4 和d=4 的输出：

tensor([[[[0.3699, 0.3584, 0.4940, 0.8618],
          [0.6767, 0.7439, 0.5984, 0.5499],
          [0.8465, 0.7276, 0.3078, 0.3882],
          [0.1001, 0.0705, 0.2007, 0.4051]]],


        [[[0.7520, 0.4528, 0.0525, 0.9253],
          [0.6946, 0.0318, 0.5650, 0.7385],
          [0.0671, 0.6493, 0.3243, 0.2383],
          [0.6119, 0.7762, 0.9687, 0.0896]]],


        [[[0.3504, 0.7431, 0.8336, 0.0336],
          [0.8208, 0.9051, 0.1681, 0.8722],
          [0.5751, 0.7903, 0.0046, 0.1471],
          [0.4875, 0.1592, 0.2783, 0.6338]]],


        [[[0.9398, 0.7589, 0.6645, 0.8017],
          [0.9469, 0.2822, 0.9042, 0.2516],
          [0.2576, 0.3852, 0.7349, 0.2806],
          [0.7062, 0.1214, 0.0922, 0.1385]]]])
tensor([[0, 3],
        [3, 2],
        [1, 1],
        [1, 0]])

我希望这是你想要的！ :)

编辑：

这是一个稍微修改过的，可能会更快（我猜不是很多：），但它更简单更漂亮：

而不是像以前那样：

m = x.view(n, -1).argmax(1)
indices = torch.cat(((m // d).view(-1, 1), (m % d).view(-1, 1)), dim=1)

已经对argmax 值进行了必要的整形：

m = x.view(n, -1).argmax(1).view(-1, 1)
indices = torch.cat((m // d, m % d), dim=1)

但正如 cmets 中所述。我认为不可能从中得到更多。

你可以做的一件事，如果真的对你来说获得最后一点可能的性能改进很重要，那就是将上面的函数实现为低级扩展（如C++) 用于 pytorch。

这只会给你一个你可以调用它的函数，并且会避免缓慢的 python 代码。

https://pytorch.org/tutorials/advanced/cpp_extension.html

【讨论】：

是的，这就是我想要的输出。我修改它以将m 转换为.float()，然后使用// 除以d。你提出的是一个解开，类似于numpy.unravel_indices()。如果你能想到一个更快的方法当然会更好。

@Chris 我刚刚做了一个短时间的测试。实际上我认为它非常有效，我想目前没有更快的方法：调用argmax() 本身需要大约 10 倍，只要计算下一行中的索引 - 在 CPU 上，我也可以稍后在 GPU 上检查。但是操作真的很简单而且很严格，所以即使这是一种解决方法，从理论的角度来看它也应该是相当有效的。

不，它一点也不慢，我在 Telsa Volta 上需要大约 5.5 毫秒。我只需要最大化它，但我同意，argmax 是一个线性运算，因为张量是无序的。可能这是最慢的组件，无法加速。

@Chris 最后我做了一个小编辑，版本稍微好一点。但我不希望在性能方面真正有所作为，可能大致相同 - 领先半纳秒。如果充分利用它真的很重要，您可能希望使用 C++ 中的自定义扩展。但考虑到代码的小sn-p，收益可能不会那么多。

谢谢，效果很好。我在评估中也犯了一个错误，似乎只是 0.5 毫秒而不是 5 毫秒。

【参考方案3】：

这是torch 中的unravel_index 实现：

def unravel_index(
    indices: torch.LongTensor,
    shape: Tuple[int, ...],
) -> torch.LongTensor:
    r"""Converts flat indices into unraveled coordinates in a target shape.

    This is a `torch` implementation of `numpy.unravel_index`.

    Args:
        indices: A tensor of (flat) indices, (*, N).
        shape: The targeted shape, (D,).

    Returns:
        The unraveled coordinates, (*, N, D).
    """

    coord = []

    for dim in reversed(shape):
        coord.append(indices % dim)
        indices = indices // dim

    coord = torch.stack(coord[::-1], dim=-1)

    return coord

然后，您可以使用torch.argmax 函数来获取“扁平化”张量的索引。

y = x.view(20, -1)
indices = torch.argmax(y)
indices.shape  # (20,)

并使用 unravel_index 函数解开索引。

indices = unravel_index(indices, x.shape[-2:])
indices.shape  # (20, 2)

【讨论】：

这是最接近真实、通用的答案！要更直接地回答原始问题，即询问如何获取最大值的索引，您可能需要编辑以显示如何首先使用 argmax 获取索引，然后解开它们。

我最终不得不对与 argmax 的连接进行编码，所以请检查我的答案。随意将我所做的纳入您的。

【参考方案4】：

接受的答案仅适用于给定的示例。

tejasvi88 的回答很有趣，但无助于回答最初的问题（正如我在那里的评论中所解释的那样）。

我相信弗朗索瓦的答案是最接近的，因为它处理的是更通用的情况（任意数量的维度）。但是，它没有与argmax 连接，并且显示的示例没有说明该函数处理批处理的能力。

因此，我将在此处以弗朗索瓦的回答为基础，并添加代码以连接到 argmax。我编写了一个新函数batch_argmax，它返回批处理中最大值的索引。批次可以按多个维度进行组织。我还包括一些测试用例来说明：

def batch_argmax(tensor, batch_dim=1):
    """
    Assumes that dimensions of tensor up to batch_dim are "batch dimensions"
    and returns the indices of the max element of each "batch row".
    More precisely, returns tensor `a` such that, for each index v of tensor.shape[:batch_dim], a[v] is
    the indices of the max element of tensor[v].
    """
    if batch_dim >= len(tensor.shape):
        raise NoArgMaxIndices()
    batch_shape = tensor.shape[:batch_dim]
    non_batch_shape = tensor.shape[batch_dim:]
    flat_non_batch_size = prod(non_batch_shape)
    tensor_with_flat_non_batch_portion = tensor.reshape(*batch_shape, flat_non_batch_size)

    dimension_of_indices = len(non_batch_shape)

    # We now have each batch row flattened in the last dimension of tensor_with_flat_non_batch_portion,
    # so we can invoke its argmax(dim=-1) method. However, that method throws an exception if the tensor
    # is empty. We cover that case first.
    if tensor_with_flat_non_batch_portion.numel() == 0:
        # If empty, either the batch dimensions or the non-batch dimensions are empty
        batch_size = prod(batch_shape)
        if batch_size == 0:  # if batch dimensions are empty
            # return empty tensor of appropriate shape
            batch_of_unraveled_indices = torch.ones(*batch_shape, dimension_of_indices).long()  # 'ones' is irrelevant as it will be empty
        else:  # non-batch dimensions are empty, so argmax indices are undefined
            raise NoArgMaxIndices()
    else:   # We actually have elements to maximize, so we search for them
        indices_of_non_batch_portion = tensor_with_flat_non_batch_portion.argmax(dim=-1)
        batch_of_unraveled_indices = unravel_indices(indices_of_non_batch_portion, non_batch_shape)

    if dimension_of_indices == 1:
        # above function makes each unraveled index of a n-D tensor a n-long tensor
        # however indices of 1D tensors are typically represented by scalars, so we squeeze them in this case.
        batch_of_unraveled_indices = batch_of_unraveled_indices.squeeze(dim=-1)
    return batch_of_unraveled_indices


class NoArgMaxIndices(BaseException):

    def __init__(self):
        super(NoArgMaxIndices, self).__init__(
            "no argmax indices: batch_argmax requires non-batch shape to be non-empty")

这里是测试：

def test_basic():
    # a simple array
    tensor = torch.tensor([0, 1, 2, 3, 4])
    batch_dim = 0
    expected = torch.tensor(4)
    run_test(tensor, batch_dim, expected)

    # making batch_dim = 1 renders the non-batch portion empty and argmax indices undefined
    tensor = torch.tensor([0, 1, 2, 3, 4])
    batch_dim = 1
    check_that_exception_is_thrown(lambda: batch_argmax(tensor, batch_dim), NoArgMaxIndices)

    # now a batch of arrays
    tensor = torch.tensor([[1, 2, 3], [6, 5, 4]])
    batch_dim = 1
    expected = torch.tensor([2, 0])
    run_test(tensor, batch_dim, expected)

    # Now we have an empty batch with non-batch 3-dim arrays' shape (the arrays are actually non-existent)
    tensor = torch.ones(0, 3)  # 'ones' is irrelevant since this is empty
    batch_dim = 1
    # empty batch of the right shape: just the batch dimension 0,since indices of arrays are scalar (0D)
    expected = torch.ones(0)
    run_test(tensor, batch_dim, expected)

    # Now we have an empty batch with non-batch matrices' shape (the matrices are actually non-existent)
    tensor = torch.ones(0, 3, 2)  # 'ones' is irrelevant since this is empty
    batch_dim = 1
    # empty batch of the right shape: the batch and two dimension for the indices since we have 2D matrices
    expected = torch.ones(0, 2)
    run_test(tensor, batch_dim, expected)

    # a batch of 2D matrices:
    tensor = torch.tensor([[[1, 2, 3], [6, 5, 4]], [[2, 3, 1], [4, 5, 6]]])
    batch_dim = 1
    expected = torch.tensor([[1, 0], [1, 2]])  # coordinates of two 6's, one in each 2D matrix
    run_test(tensor, batch_dim, expected)

    # same as before, but testing that batch_dim supports negative values
    tensor = torch.tensor([[[1, 2, 3], [6, 5, 4]], [[2, 3, 1], [4, 5, 6]]])
    batch_dim = -2
    expected = torch.tensor([[1, 0], [1, 2]])
    run_test(tensor, batch_dim, expected)

    # Same data, but a 2-dimensional batch of 1D arrays!
    tensor = torch.tensor([[[1, 2, 3], [6, 5, 4]], [[2, 3, 1], [4, 5, 6]]])
    batch_dim = 2
    expected = torch.tensor([[2, 0], [1, 2]])  # coordinates of 3, 6, 3, and 6
    run_test(tensor, batch_dim, expected)

    # same as before, but testing that batch_dim supports negative values
    tensor = torch.tensor([[[1, 2, 3], [6, 5, 4]], [[2, 3, 1], [4, 5, 6]]])
    batch_dim = -1
    expected = torch.tensor([[2, 0], [1, 2]])
    run_test(tensor, batch_dim, expected)


def run_test(tensor, batch_dim, expected):
    actual = batch_argmax(tensor, batch_dim)
    print(f"batch_argmax of tensor with batch_dim batch_dim is\nactual\nExpected:\nexpected")
    assert actual.shape == expected.shape
    assert actual.eq(expected).all()

def check_that_exception_is_thrown(thunk, exception_type):
    if isinstance(exception_type, BaseException):
        raise Exception(f"check_that_exception_is_thrown received an exception instance rather than an exception type: "
                        f"exception_type")
    try:
        thunk()
        raise AssertionError(f"Should have thrown exception_type")
    except exception_type:
        pass
    except Exception as e:
        raise AssertionError(f"Should have thrown exception_type but instead threw e")

【参考方案5】：

ps=ps.numpy()
ps=ps.tolist()

mx=[max(l) for l in ps]
mx=max(mx)
for i in range(len(ps[0])):
  if mx==ps[0][i]:
    print("The digit is "+str(i))
    break

这对我很有效