如何在numpy或pytorch中向量化自定义算法? [英] How to vectorize custom algorithms in numpy or pytorch?

问题描述

假设我有两个矩阵:

A: size k x m

B: size m x n

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

Using a custom operation, my output will be k x n.

此自定义操作不是A行和B列之间的点积. 假设，此自定义操作定义为:

This custom operation is not a dot product between the rows of A and columns of B. Suppose this custom operation is defined as:

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

For the Ith row of A and Jth column of B, the i,j element of the output is:

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

我看到的唯一实现此方法的方法是扩展此方程式，计算每个项，然后将它们求和.

The only way I can see to implement this is to expand this equation, calculate each term, them sum them.

在numpy或pytorch中是否有一种方法可以在不扩展方程式的情况下进行操作?

Is there a way in numpy or pytorch to do this without expanding the equation?

推荐答案

除了注释中的@hpaulj方法外，您还可以使用以下事实:您要计算的本质上是成对的Minkowski距离:

Apart from the method @hpaulj outlines in the comments, you can also use the fact that what you are calculating is essentially a pair-wise Minkowski distance:

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

这篇关于如何在numpy或pytorch中向量化自定义算法?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！