我需要乘以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的机器。在
Tags:
根据您的矩阵
C
,稀疏矩阵可能会减少大小和计算时间,例如,您只保存非0的列,torch reference可能会有所帮助。在因为你有四个gpu,你可以利用它们来执行高效的矩阵乘法。但是请注意,乘法的结果大小为3000x360000,在单精度浮点(fp32)中占40GB。除非有足够大的内存供CPU使用,否则无法将计算结果存储在RAM上。在
一个可能的解决方案是将大矩阵
C
分成四个较小的块,在不同的GPU上执行每个块的矩阵乘法,并将结果保存在GPU上。如果每个GPU至少有10GB的内存,那么您将有足够的内存来完成此任务。在如果你也有足够的CPU内存,那么你可以将所有四个GPU的结果移动到CPU上并将它们连接起来(实际上,在这种情况下,您可以只使用一个GPU,并且每次都将结果从GPU传输到CPU)。否则,您可以将结果分块保存在gpu上,您需要记住并跟踪这四个块实际上是一个矩阵的一部分。在
如果有多个gpu,可以使用PyTorch的^{} 在所有gpu上分布计算。它将在gpu中拆分(并行化)矩阵
C_gpu
列的乘法。在方法如下:
首先,导入模块并准备矩阵:
下一步,创建一个没有偏移的^{} “层”。这个层所做的就是矩阵乘法。输入大小将是
^{pr2}$C_gpu
中每列的大小,而输出大小将是结果中每列的大小。在将层的矩阵(=weight)设置为
A_gpu @ B_gpu
,这是一个可以快速计算而不需要并行化的小矩阵(尽管您也可以根据需要并行化它)。在将层转换为DataParallel实例。这意味着它将沿着“批处理”维度自动并行计算。参数
device_ids
是gpu的索引列表(在您的例子中,其中有4个)。在现在您可以将矩阵
C_gpu
输入到层中,计算将沿着其大维度并行进行:重要提示:在编写此答案时,我无法访问多个GPU来实际测试此建议的解决方案。如果有读者能确认它确实有效,我将不胜感激(甚至可以更好地-计时解决方案并与单个GPU进行比较)
EDIT1:我现在在多个gpu(四个Nvidia Tesla P100)上尝试了这个代码,结果发现它给出了内存不足的错误。我将把这个解决方案作为参考,因为它确实适用于尺寸高达400K(而不是3.6M)。在
另外,如果您将
C
分成更小的块,将每个块输入mat_mult_gpu
,然后将结果连接到CPU上,那么这个解决方案仍然适用于3.6M的大小。注意,您需要大量的CPU内存才能工作,因为结果的大小是3K-by-3.6M
,在fp32中大约需要40GB。(或者,您可以将每个块保存到磁盘,而无需连接块)。在相关问题 更多 >
编程相关推荐