如何在 numpy 或 pytorch 中向量化自定义算法？

Question

假设我有两个矩阵：

A: size k x m

B: size m x n

使用自定义操作，我的输出将是 k x n.

此自定义操作不是 A 的行和 B 的列之间的点积。假设这个自定义操作定义为：

对于A的第I行和B的第J列，输出的i,j元素是：

sum( (a[i] + b[j]) ^20 ), i loop over I, j loops over J

我能看到实现这个的唯一方法是扩展这个方程，计算每一项，然后将它们相加。

在 numpy 或 pytorch 中有没有办法在不扩展方程的情况下做到这一点？

Answer 1

除了@hpaulj 在评论中概述的方法之外，您还可以使用这样一个事实，即您正在计算的本质上是成对的 Minkowski 距离：

import numpy as np
from scipy.spatial.distance import cdist

k,m,n = 10,20,30
A = np.random.random((k,m))
B = np.random.random((m,n))

method1 = ((A[...,None]+B)**20).sum(axis=1)
method2 = cdist(A,-B.T,'m',p=20)**20

np.allclose(method1,method2)
# True

Answer 2

你可以自己实现

以下函数生成各种类似函数的点积，但不要用它来代替 np.dot，因为对于较大的数组，它会慢很多。

模板

import numpy as np
import numba as nb
from scipy.spatial.distance import cdist

def gen_dot_like_func(kernel,parallel=True):

    kernel_nb=nb.njit(kernel,fastmath=True)

    def cust_dot(A,B_in):
        B=np.ascontiguousarray(B_in.T)
        assert B.shape[1]==A.shape[1]

        out=np.empty((A.shape[0],B.shape[0]),dtype=A.dtype)
        for i in nb.prange(A.shape[0]):
            for j in range(B.shape[0]):
                sum=0
                for k in range(A.shape[1]):
                    sum+=kernel_nb(A[i,k],B[j,k])
                out[i,j]=sum
        return out

    if parallel==True:
        return nb.njit(cust_dot,fastmath=True,parallel=True)
    else:
        return nb.njit(cust_dot,fastmath=True,parallel=False)

生成函数

#This can be useful if you have a lot matrix-multiplication like functions
my_func=gen_dot_like_func(lambda A,B:(A+B)**20,parallel=True)

时间

k,m,n = 10,20,30
%timeit method1 = ((A[...,None]+B)**20).sum(axis=1)
192 µs ± 554 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit method2 = cdist(A,-B.T,'m',p=20)**20
208 µs ± 1.85 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit res=my_func(A,B) #parallel=False
4.01 µs ± 34.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

k,m,n = 500,100,500
timeit method1 = ((A[...,None]+B)**20).sum(axis=1)
852 ms ± 4.93 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit method2 = cdist(A,-B.T,'m',p=20)**20
714 ms ± 2.12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res=my_func(A,B) #parallel=True
1.81 ms ± 11.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

如何在 numpy 或 pytorch 中向量化自定义算法？

How to vectorize custom algorithms in numpy or pytorch?

python

numpy

vectorization

pytorch

你可以自己实现