Python 如何将张量分解成小块进行矩阵乘法以解决内存问题

我有两个矩阵： 火炬尺寸（[1,8,50176,64]）火炬尺寸（[1,8,64,50176]） 我需要执行矩阵乘法来获得torch大小[1,8,50176,50176]，但我的cuda内存不足，需要75GB的GPU内存，我负担不起那么多GPU，我的GPU是8GB，我想知道如何将矩阵分解成小块并执行乘法，最后在一个torch张量中得到结果

我试过：

y = torch.matmul(a, b) 
y = torch.bmm(a.view(-1, 50176, 64), b.view(-1, 64, 50176)).view(1, 8, 50176, 50176)
y = torch.einsum('bdij, bdjk -> bdik', (a, b))

以上工作均未完成，均导致cuda内存不足错误

谢谢你的帮助

---

Ling

A（float）形状张量（1,85017650176）占据了4*（8*50176*50176）/2^30=75GB的空间。将其拆分为更小的矩阵乘法并不能解决这个问题，您仍然需要75GB来存储最终矩阵。正如jodag所说，最终张量大小存在瓶颈。如果可能（精度方面）将

a

和

b

转换为

torch.cuda.HalfTensor

，然后每次计算大小为1的批次的乘法，这应该需要4.6875GB的VRAM。然后，您有两个选项：如果您有大量RAM，则在结果矩阵上使用

.to（'cpu'）

将张量推入，最后在批处理维度上连接所有结果矩阵。如果您没有太多RAM，请将每个乘法结果保存在磁盘上，并处理数据的加载（/卸载）。谢谢你们的评论，我有一台16G RAM（仅CPU）的笔记本电脑，它可以用[900,8,64,50176]计算[900,8,50176,64]的matmul，而没有任何问题，但GPU机器无法做到，我把GPU机器换成了拥有32G内存的CPU，但仍然无法做到，为什么16G内存可以做到呢？@Ling CPU内存可以是“虚拟的”：操作系统可以将部分RAM内存交换到磁盘上，即使您消耗的内存超过了物理RAM的实际大小，也可以让您继续运行。操作系统无法对GPU内存执行此操作。