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() 并使用 30sample 时间窗口（即它只返回 60 个时间点窗口内的最大值）。

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

Answer 1

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

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

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

Answer 2

好的，所以 pytorch 论坛上的某个人确实有一个很好的解决方案： https ://discuss.pytorch.org/t/pytorch-argrelmax-or-c-function/36404

这是完整的答案，以防它被删除：

    a = #1D Array of  your choice
    window_maxima = torch.nn.functional.max_pool1d_with_indices(a.view(1,1,-1), width, 1, padding=width//2)[1].squeeze()
    candidates = window_maxima.unique()
    nice_peaks = candidates[(window_maxima[candidates]==candidates).nonzero()]