计算矩阵乘法的最快方法

Question

我按如下方式计算矩阵乘法：

    import numpy as np
    A=np.random.randn(1,3000,30000)
    B=np.random.randn(50,1,30000)
    C=A*B

C的计算大约需要10分钟。我怎么能改善这个？

    In [4]: timeit A*B
    1 loops, best of 3: 1min 59s per loop

更新

Answer 1

看来你的主要瓶颈是RAM大小。计算结果 36 GB ，使您的操作系统将数据写入交换。您可以尝试使用dask进行核心计算，并直接将结果传输到HDF5文件：

import dask
import dask.array

A = dask.array.random.random((1, 3000, 300000), chunks=1024)
B = dask.array.random.random((50, 1, 300000), chunks=1024)

C = A * B

dask.array.to_hdf5('myfile.hdf5', '/C', C)

确保安装

pip install dask[array]
pip instal h5py

这可能不会加速您的计算，因为您仍然受I / O限制（仍然写入硬盘），但至少它会使计算可管理并使您的计算机保持响应。

另一种解决方案是手动切片操作，如果您不需要全部C以供将来计算：

for i in B.shape[0]:
    C = A * B[i, ...]  # do not save this result but rather use and discard it immediately