如何在 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屋！