python - Pytorch Argrelmax 函数(或 C++)

Question

我正在尝试为 scipy.signal.argrelmax() 找到等效的 pytorch(或 C++)，它在带有一些填充的一维数组中找到峰值。 https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.signal.argrelmax.html

这是我想出的，它比 scipy.signal.argrelmax 更快 - 但我缺少对删除某个窗口内的峰值的最后一步的快速解决方案。

import torch

# initalize an array (not the one in plot below)
gpu_max = torch.rand(100000)

# find peaks and troughs by subtracting shifted versions
gpu_temp1 = gpu_max[1:-1]-gpu_max[:-2]
gpu_temp2 = gpu_max[1:-1]-gpu_max[2:]

# and checking where both shifts are positive;
out1 = torch.where(gpu_temp1>0, gpu_temp1*0+1, gpu_temp1*0)
out2 = torch.where(gpu_temp2>0, out1, gpu_temp2*0)

# argrelmax containing all peaks
argrelmax_gpu = torch.nonzero(out2, out=None)+1

所以标记每个 relmaxpeak 的顶部图非常快。但需要底部的 - 它来自 scipy.signal.argrelmax() 并使用 30 个样本时间窗口(即它仅在 60 个时间点窗口内返回最大值)。

编辑:我更新了代码以反射(reflect)更新的 torch.nonzero() 搜索。但仍然需要弄清楚如何做最后一步(如下图)。

Answer 1

PyTorch 中没有这样的东西，我想不出一种简单的方法来根据现成的函数高效地实现它，抱歉。用 C++ 编写代码使其正常工作非常简单，使其与大颗粒(想想 CPU)并行运行也相当简单，使其在具有 1000 多个内核的 GPU 上运行则更加棘手，因为您必须在并行性之间找到正确的平衡以及后续索引计算之间的数据重用。尽管如此，如果您决定采用这种方式，tutorial在自定义 native 扩展上非常好。

至于 torch.where 问题，此功能是通过 torch.nonzero 提供的.

编辑:请注意，您上面的方法可以通过在 for 循环中迭代扩展到任何大小的填充，并且它的算法性能不会低于 scipy 东西，它也做了一个非常简单的 argmax 窗口应用程序，而不是迭代一些聪明的优先级列表。因此，在 GPU 支持下，您可能仍会看到可接受的性能(取决于您的用例)。