如何在pytorch中批量进行矩阵向量乘法（一个矩阵，多个向量）而不复制内存中的矩阵

Question

我有

n

大小为

d

的向量和单个

d x d

矩阵

J

。我想计算

n

与每个

J

向量的

n

矩阵向量乘法。

为此，我使用 pytorch 的

expand()

来获得 J 的

broadcast

，但似乎在计算矩阵向量乘积时，pytorch 在内存中实例化了完整的

n x d x d

张量。例如以下代码

device = torch.device("cuda:0")
n = 100_000_000
d = 10

x = torch.randn(n, d, dtype=torch.float32, device=device)
J = torch.randn(d, d, dtype=torch.float32, device=device).expand(n, d, d)
y = torch.sign(torch.matmul(J, x[..., None])[..., 0])

加薪

RuntimeError: CUDA out of memory. Tried to allocate 37.25 GiB (GPU 0; 11.00 GiB total capacity; 3.73 GiB already allocated; 5.69 GiB free; 3.73 GiB reserved in total by PyTorch)

这意味着 pytorch 不必要地尝试为矩阵的

n

副本分配空间

J

如何以矢量化方式执行此任务（矩阵很小，因此我不想循环每个矩阵向量乘法）而不耗尽 GPU 内存？

Answer 1

我想这会解决这个问题：

import torch
x = torch.randn(n, d)
J = torch.randn(d, d) # no need to expand

y = torch.matmul(J, x.T).T

用你的表情验证：

Jex = J.expand(n, d, d)
y1 = torch.matmul(Jex, x[..., None])[..., 0]
y = torch.matmul(J, x.T).T

torch.allclose(y1, y) # using allclose for float values
# tensor(True)

Answer 2

我相信这是一个解决方案：

进口火炬

x = torch.randn(n, d)

J = torch.randn(d, d)

y = torch.einsum('ij,bj->bi', J, x) # y.shape -> (n,d)