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

Question

我有n 大小为d 的向量和一个d x d 矩阵J。我想计算J 与每个n 向量的n 矩阵向量乘法。

为此，我使用 pytorch 的expand() 来获得J 的广播，但似乎在计算矩阵向量积时，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 不必要地尝试为矩阵 J 的 n 副本分配空间

如何在不耗尽 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)