CUDA内存与矩阵相乘

Question

我试图将3个矩阵相乘，但是CUDA内存耗尽了。

# A: 3000 x 100 (~2MB)
# B: 100 x 100  (~0.05MB)
# C: 100 x 3MM  (~2GB)

A = np.random.randn(3000, 100)
B = np.random.randn(100, 100)
C = np.random.randn(100, 3e6)

A_gpu = torch.from_numpy(A).cuda()
B_gpu = torch.from_numpy(B).cuda()
C_gpu = torch.from_numpy(C).cuda()

R_gpu = (A_gpu @ B_gpu @ C_gpu)

Cuda正在为此操作请求大约90GB的内存。我不明白为什么。

Answer 1

乘以矩阵，您的输出大小将为3,000 x 3,000,000矩阵！所以尽管A和B相对较小，但输出R是巨大的：9G元素。此外，我怀疑你的矩阵的dtype是float64而不是float32（因为你使用numpy来启动它们）。因此，R_gpu的每个9G元素需要8个字节;为R_gpu带来至少72GB GPU内存的大小。我怀疑中间结果和其他一些东西占用了你的GPU内存。