如何有效地乘以重复行的火炬张量而不将所有行存储在内存中或迭代？

Posted 2023-03-29

【中文标题】如何有效地乘以重复行的火炬张量而不将所有行存储在内存中或迭代？

【英文标题】：How to efficiently multiply by torch tensor with repeated rows without storing all the rows in memory or iterating?

【发布时间】：2021-08-02 09:13:34

【问题描述】：

给定一个火炬张量：

# example tensor size 2 x 4
a = torch.Tensor([[1, 2, 3, 4], [5, 6, 7, 8]])

还有一个每 n 行重复一次的地方：

# example tensor size 4 x 3 where every 2 rows repeated
b = torch.Tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]])

如何进行矩阵乘法：

>>> torch.mm(a, b)
tensor([[ 28.,  38.,  48.],
        [ 68.,  94., 120.]])

不将整个重复的行张量复制到内存中或迭代？

即只存储前 2 行：

# example tensor size 2 x 3 where only the first two rows from b are actually stored in memory
b_abbreviated = torch.Tensor([[1, 2, 3], [4, 5, 6]])

因为这些行将被重复。

有一个功能

torch.expand()

但这在重复多行时确实有效，并且作为这个问题：

Repeating a pytorch tensor without copying memory

表明并且我自己的测试确认在调用时通常最终会将整个张量复制到内存中

.to(device)

也可以迭代地执行此操作，但速度相对较慢。

有没有什么方法可以有效地执行此操作，而不会将整个重复的行张量存储在内存中？

编辑说明：

对不起，最初没有澄清：一个被用作第一个张量的第一个维度以保持示例简单，但我实际上正在寻找任何两个张量 a 和 b 的一般情况的解决方案，使得它们的维度与矩阵乘法兼容，并且 b 的行每 n 行重复一次。我更新了示例以反映这一点。

【参考方案1】：

假设a 的第一个维度为 1，如您的示例所示，您可以执行以下操作：

a = torch.Tensor([[1, 2, 3, 4]])
b_abbreviated = torch.Tensor([[1, 2, 3], [4, 5, 6]])
torch.mm(a.reshape(-1, 2), b_abbreviated).sum(axis=0, keepdim=True)

在这里，不是重复行，而是将a 乘以块，然后将它们逐列相加以获得相同的结果。

---

如果a的第一个维度不一定是1，你可以尝试如下：

torch.cat(torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0]), dim=1).sum(
dim=0, keepdim=True).reshape(a.shape[0], -1)

在这里，您执行以下操作：

使用torch.mm(a.reshape(-1,2),b_abbreviated，您再次将a 的每一行拆分为大小为2 的块，并将它们一个接一个地堆叠起来，然后将每一行一个接一个地堆叠起来。 使用torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0])，这些堆栈然后按行分开，因此拆分的每个结果组件对应于单行的块。 使用torch.cat(torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0]), dim=1)，这些堆栈然后按列连接。 使用.sum(dim=0, keepdim=True)，将对应于a 中各个行的单独块的结果相加。 使用.reshape(a.shape[0], -1)，按列连接的a 行再次按行堆叠。

与直接矩阵乘法相比，它似乎相当慢，这并不奇怪，但与显式迭代相比，我还没有检查过。可能有更好的方法可以做到这一点，如果我想到任何会编辑。

【讨论】：

对不起，对于我最初给出的例子来说，这看起来是一个很好的答案，但我实际上并不是要暗示第一个张量的第一维总是一个。我现在用一个更一般的例子更新了这个问题来澄清。

啊，它可以被概括，这很酷。但是，我注意到由 torch.cat(torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0]), dim=1 生成的中间矩阵的大小) 的大小与 a 的第一个维度成比例地增加，这最终可能会破坏当 a 具有大量行时试图避免使用大量内存的目的。我很想知道这样大的中间矩阵是否可以避免，但我认为这可能是解决这个问题的最佳方法。

@nellapizza 中间矩阵的大小只有最终矩阵的 2 倍。我假设在b 中的重复次数非常高的情况下会节省任何费用。无论如何，就时间效率而言，我认为这根本不是最好的方法，但我想不出更好的方法。