如何并行化 for 循环以在 PyTorch 中使用？答案

【问题标题】：How can I parallelize a for loop for use in PyTorch?如何并行化 for 循环以在 PyTorch 中使用？
【发布时间】：2020-07-20 03:57:57
【问题描述】：

我意识到for 循环通常与Python 相比很慢。我有一些代码弄乱了一些张量：


            for batch_index, mask_batch in enumerate(mask):
                mask_len = torch.sum(mask_batch).int()

                if mask_len == 0:
                    side_input = torch.zeros((max_inp_len, side_input.shape[1])).to(mask.device)
                else:

                    m_nonzero = mask_batch.nonzero().flatten()
                    first_nonzero = m_nonzero[0]
                    last_nonzero = m_nonzero[-1]

                    if side == 'left':
                        end_index = first_nonzero - 1
                        start_index = 0
                    elif side == 'right':
                        start_index = last_nonzero + 1
                        end_index = inputs[batch_index].size(1)

                    side_input = inputs[batch_index][start_index:end_index]

                    if end_index - start_index < max_inp_len:
                        pad_zeros = torch.zeros(
                            (max_inp_len - side_input.shape[0], side_input.shape[1])).to(mask.device)
                        if side == 'left':
                            side_input = torch.cat((pad_zeros, side_input), 0)
                        elif side == 'right':
                            side_input = torch.cat((side_input, pad_zeros), 0)

                side_inputs.append(side_input)

        return torch.stack(side_inputs)

我觉得这个循环真的让事情变慢了。有什么方法可以让我在没有循环的情况下做到这一点？

【问题讨论】：

for 循环不是“慢”吗……你怎么能这么说？无论如何，如果您在循环中执行许多慢速操作，请考虑使用 threading 或 multiprocessing 库。
如果你的 for 循环包括遍历一个张量，并对张量的元素执行一些操作，你应该尝试向量化这些操作。
会不会因为你在累积梯度而变慢？如果你把with torch.no_grad()放在你所有的代码前面呢？

标签： python arrays numpy pytorch tensor

【解决方案1】：

Python 在任何给定进程中都没有真正的并行性。您必须生成一个 ProcessPool 并使循环内部成为一个采用 batch_index, mask_batch 的函数，然后将该函数映射到当前 for 循环中的 mask 对象。问题是，我不知道 PyTorch 是否能很好地处理这个问题。

像这样

def f(batch_index, mask_batch):
    mask_len = torch.sum(mask_batch).int()

    if mask_len == 0:
        side_input = torch.zeros((max_inp_len, side_input.shape[1])).to(mask.device)
    else:
        m_nonzero = mask_batch.nonzero().flatten()
        first_nonzero = m_nonzero[0]
        last_nonzero = m_nonzero[-1]

        if side == 'left':
            end_index = first_nonzero - 1
            start_index = 0
        elif side == 'right':
            start_index = last_nonzero + 1
            end_index = inputs[batch_index].size(1)

            side_input = inputs[batch_index][start_index:end_index]

            if end_index - start_index < max_inp_len:
                pad_zeros = torch.zeros((max_inp_len - side_input.shape[0], side_input.shape[1])).to(mask.device)
                if side == 'left':
                    side_input = torch.cat((pad_zeros, side_input), 0)
                elif side == 'right':
                    side_input = torch.cat((side_input, pad_zeros), 0)
    return side_input

您可以查看的其他内容是进一步矢量化代码。 PyTorch 和 Numpy 中的大多数东西都可以通过使用内置函数并将另一个维度添加到表示“循环”维度的张量上来向量化。这将允许 PyTorch 为您处理并行性。

PyTorch 可能有一个设备概念，您可以在其中放置循环的不同迭代，同样这将需要您为此循环创建一个函数，并可能将其继续运行的设备作为输入。

最后，您可以研究像 Numba 或 torch.jit 这样的即时编译来为您执行自动矢量化。

如果mask 的长度未知，这一切都不起作用（很可能）。如果它的长度是已知的，我认为尽管矢量化很困难，但它可能是您的最佳选择。

【讨论】：

【解决方案2】：

您应该创建一个包含循环迭代背后的逻辑的函数，并将其作为每个列的线程启动（请参阅docs here）。您也可以使用 asyncio 库进行并发，但您可能会获得较少的改进。

为列表的每个元素生成线程的一个很好的例子可以阅读here。

【讨论】：