PyTorch For 循环优化和加速技术

Question

这是我在过去一年中遇到过三次的问题。

我很欣赏，在某些情况下，矢量化解决方案会更好并且速度更快。

但是恕我直言，在发现矢量化解决方案和使用实质上的 for 循环（或双 for 循环）之间存在权衡。发现矢量化解决方案（如果确实存在）可能需要付出更多的努力、反复试验、研究等等。

最简单形式的代码（在本例中为双 for 循环）几乎总是成为我的瓶颈，但只需很少的时间来实现和测试。

这是一个例子：

@torch.jit.script    
def seq_prob(t_samples: torch.Tensor):
    i = 0
    probs = [0] * len(t_samples)    
    for t_i in torch.unbind(t_samples):
        for t_k in torch.unbind(t_samples):
            is_same = torch.all(torch.isclose(t_i, t_k, rtol=1e-05, atol=1e-08, equal_nan=False))
            if is_same is True:
                probs[i] += 1
        i += 1
    return probs

torch.unbind

只是将外部维度视为可迭代。在某些情况下，我花了相当多的时间来导出循环的向量化形式，这通常会导致屏蔽、求和、索引选择和各种内置的 pytorch 方法，与 for 循环相比，使逻辑变得复杂但让它更快。

同样，使用 CUDA 和

@torch.jit

有时会有所帮助（但并非总是如此）。

因此我的问题是：

• 在 pytorch 中使用某种形式的 for 循环时（例如，

  torch.unbind

  或

  torch.chunk

  或任何类似的东西），其目的是迭代维度并执行操作
• 有没有一种方法，一个黄金标准，一些选项，可以加快速度（矢量化除外）？
• 如果矢量化是唯一的选择，那么第一个好的攻击计划是什么？以上面显示的代码为例，它计算给定一定容差的样本集中的值的观测值。

Answer 1

我无法直接提出矢量化版本，但我修复了一个错误并将循环运行减少了 50% 以上。

我无法为您提供黄金标准，除了：尽可能矢量化，当可以避免时不要在两个方向上进行成对比较。

import torch

t = torch.Tensor([[1, 2, 3], [1, 2, 3], [1, 1, 1]])
torch.all(torch.isclose(t[0], t[0]))


t_samples = t
i = 0
probs = torch.zeros(len(t_samples))
for id_i, t_i in enumerate(torch.unbind(t_samples)):
    # dont do the same calculation twicse, start at id_i + 1
    for id_j, t_k in enumerate(torch.unbind(t_samples[id_i+1:]), start=id_i+1):
        is_same = torch.all(
            torch.isclose(t_i, t_k, rtol=1e-05, atol=1e-08, equal_nan=False)
        )
        if is_same:  # you had an error here, torch booleans dont work your way
            probs[id_i] += 1  # compare A to B
            probs[id_j] += 1  # compare B to A