一个矩阵极宽时实现矩阵乘法的有效方法？

Question

我需要乘以3个矩阵，A: 3000x100, B: 100x100, C: 100x3.6MM。我现在只是在PyTorch中使用普通矩阵乘法

A_gpu = torch.from_numpy(A)
B_gpu = torch.from_numpy(B)
C_gpu = torch.from_numpy(C)
D_gpu = (A_gpu @ B_gpu @ C_gpu.t()).t()

C语言非常广泛，所以gpu上的数据重用是有限的，但是有没有其他方法可以加快这个速度？我有一台4倍GPU的机器。在

Answer 1

根据您的矩阵C，稀疏矩阵可能会减少大小和计算时间，例如，您只保存非0的列，torch reference可能会有所帮助。在

Answer 2

因为你有四个gpu，你可以利用它们来执行高效的矩阵乘法。但是请注意，乘法的结果大小为3000x360000，在单精度浮点（fp32）中占40GB。除非有足够大的内存供CPU使用，否则无法将计算结果存储在RAM上。在

一个可能的解决方案是将大矩阵C分成四个较小的块，在不同的GPU上执行每个块的矩阵乘法，并将结果保存在GPU上。如果每个GPU至少有10GB的内存，那么您将有足够的内存来完成此任务。在

如果你也有足够的CPU内存，那么你可以将所有四个GPU的结果移动到CPU上并将它们连接起来（实际上，在这种情况下，您可以只使用一个GPU，并且每次都将结果从GPU传输到CPU）。否则，您可以将结果分块保存在gpu上，您需要记住并跟踪这四个块实际上是一个矩阵的一部分。在

import numpy as np
import torch.nn as nn
import torch

number_of_gpus = 4

# create four matrics
A = np.random.normal(size=(3000,100))
B = np.random.normal(size=(100,100))
C = np.random.normal(size=(100,3600000))

# convert them to pytorch fp32 tensors
A = torch.from_numpy(A).float()
B = torch.from_numpy(B).float()
C = torch.from_numpy(C).float()

# calcualte `A@B`, which is easy
AB = A@B

# split the large matrix `C` into 4 smaller chunks along the second dimension. 
# we assume here that the size of the second dimension of `C` is divisible by 4.  
C_split = torch.split(C,C.shape[1]//number_of_gpus,dim=1)

# loop over the four GPUs, and perform the calculation on each using the corresponding chunk of `C`
D_split = []
for i in range(number_of_gpus):
    device = 'cuda:{:d}'.format(i)
    D_split.append( AB.to(device) @ C_split[i].to(device))

# DO THIS ONLY IF YOU HAVE ENOUGH CPU MEMORY!! :
D = torch.cat([d.cpu() for d in D_split],dim=1)

Answer 3

如果有多个gpu，可以使用PyTorch的^{}在所有gpu上分布计算。它将在gpu中拆分（并行化）矩阵C_gpu列的乘法。在

方法如下：

首先，导入模块并准备矩阵：

import torch
import torch.nn as nn

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

下一步，创建一个没有偏移的^{}“层”。这个层所做的就是矩阵乘法。输入大小将是C_gpu中每列的大小，而输出大小将是结果中每列的大小。在

^{pr2}$

将层的矩阵（=weight）设置为A_gpu @ B_gpu，这是一个可以快速计算而不需要并行化的小矩阵（尽管您也可以根据需要并行化它）。在

mat_mult.weight.data = A_gpu @ B_gpu

将层转换为DataParallel实例。这意味着它将沿着“批处理”维度自动并行计算。参数device_ids是gpu的索引列表（在您的例子中，其中有4个）。在

mat_mult_gpu = nn.DataParallel(mat_mult,device_ids=[0,1,2,3]).to('cuda:0')

现在您可以将矩阵C_gpu输入到层中，计算将沿着其大维度并行进行：

D_gpu  = mat_mult_gpu(C_gpu.t())

---

重要提示：在编写此答案时，我无法访问多个GPU来实际测试此建议的解决方案。如果有读者能确认它确实有效，我将不胜感激（甚至可以更好地-计时解决方案并与单个GPU进行比较）

---

EDIT1:我现在在多个gpu（四个Nvidia Tesla P100）上尝试了这个代码，结果发现它给出了内存不足的错误。我将把这个解决方案作为参考，因为它确实适用于尺寸高达400K（而不是3.6M）。在

另外，如果您将C分成更小的块，将每个块输入mat_mult_gpu，然后将结果连接到CPU上，那么这个解决方案仍然适用于3.6M的大小。注意，您需要大量的CPU内存才能工作，因为结果的大小是3K-by-3.6M，在fp32中大约需要40GB。（或者，您可以将每个块保存到磁盘，而无需连接块）。在