时间戳

首页 > 解决方案 > 如何在 numpy 或 pytorch 中矢量化自定义算法？

python - 如何在 numpy 或 pytorch 中矢量化自定义算法？

问题描述

假设我有两个矩阵：

A: size k x m

B: size m x n

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

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

对于 的第 I 行第AJ 列B，输出的i,j元素为：

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

我能看到实现这一点的唯一方法是扩展这个方程，计算每个项，然后对它们求和。

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

标签： pythonnumpyvectorizationpytorch

解决方案

除了@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

推荐阅读