torch.topk的时间复杂度

Question

当n.topk(m)的m和n分别超过2千万和200,000时，排序变得非常慢(超过3小时)。我想知道torch.topk的时间复杂度和排序的改进措施。

topv, topi = outline.topk(beam_size)  # beam_size = 200,000, outline: 1 × 20,000,000

Answer 1

我不知道pytorch是如何为CPU张量实现topk的。但是，由于您使用的是CPU，因此可以使用numpy arrays的现有部分排序实现。

例如，使用bottleneck.argpartition

import bottleneck
with torch.no_grad():
  topi = bottleneck.argpartition(outline.numpy(), kth=beam_size)
topv = outline[topi]  # allow gradients to propagate through indexing

请注意，bottleneck.argpartition的有效实现不会对数组进行排序，因此可以保证topv确实大于数组的所有其他元素，但是topv中的值不会排序，对于某些i，可能会出现topv[i] < topv[i+1]的情况。