numpy 如何让python中的最小加矩阵乘法更快?

p1iqtdky  于 2022-12-23  发布在  Python
关注(0)|答案(3)|浏览(232)

我有两个矩阵A和B,我想计算最小加乘积,如下所示:Min-plus matrix multiplication。为此,我实现了以下代码:

def min_plus_product(A,B):
    B = np.transpose(B)
    Y = np.zeros((len(B),len(A)))
    for i in range(len(B)):
         Y[i] = (A + B[i]).min(1)
    return np.transpose(Y)

这个方法很好用,但是对于大矩阵来说速度很慢,有没有办法让它更快?我听说用C实现或者使用GPU可能是不错的选择。

ddrv8njm

ddrv8njm1#

如果中间维足够大,并且元素均匀分布,那么下面的算法可以节省一点,它利用了最小和通常来自两个小项的事实。

import numpy as np

def min_plus_product(A,B):
    B = np.transpose(B)
    Y = np.zeros((len(B),len(A)))
    for i in range(len(B)):
         Y[i] = (A + B[i]).min(1)
    return np.transpose(Y)

def min_plus_product_opt(A,B, chop=None):
    if chop is None:
        # not sure this is optimal
        chop = int(np.ceil(np.sqrt(A.shape[1])))
    B = np.transpose(B)
    Amin = A.min(1)
    Y = np.zeros((len(B),len(A)))
    for i in range(len(B)):
        o = np.argsort(B[i])
        Y[i] = (A[:, o[:chop]] + B[i, o[:chop]]).min(1)
        if chop < len(o):
            idx = np.where(Amin + B[i, o[chop]] < Y[i])[0]
            for j in range(chop, len(o), chop):
                if len(idx) == 0:
                    break
                x, y = np.ix_(idx, o[j : j + chop])
                slmin = (A[x, y] + B[i, o[j : j + chop]]).min(1)
                slmin = np.minimum(Y[i, idx], slmin)
                Y[i, idx] = slmin
                nidx = np.where(Amin[idx] + B[i, o[j + chop]] < Y[i, idx])[0]
                idx = idx[nidx]
    return np.transpose(Y)

A = np.random.random(size=(1000,1000))
B = np.random.random(size=(1000,2000))

print(np.allclose(min_plus_product(A,B), min_plus_product_opt(A,B)))

import time
t = time.time();min_plus_product(A,B);print('naive {}sec'.format(time.time()-t))
t = time.time();min_plus_product_opt(A,B);print('opt {}sec'.format(time.time()-t))

输出示例:

True
naive 7.794037580490112sec
opt 1.65810227394104sec
y0u0uwnf

y0u0uwnf2#

一种可能的简单方法是使用numba

from numba import autojit
import numpy as np
@autojit(nopython=True)
def min_plus_product(A,B):
    n = A.shape[0]
    C = np.zeros((n,n))
    for i in range(n):
        for j in range(n):
            minimum = A[i,0]+B[0,j]
            for k in range(1,n):
                minimum = min(A[i,k]+B[k,j],minimum)
            C[i,j] = minimum
    return C

1000x1000 A、B矩阵上的时序为:
1圈,3局两胜:原始代码的每个循环4.28 s
1圈,3局两胜:数字代码每循环2.32 s

crcmnpdw

crcmnpdw3#

下面是一个简洁而完全 numpy 解决方案,没有任何基于Python的循环:

(np.expand_dims(a, 0) + np.expand_dims(b.T, 1)).min(axis=2).T

相关问题