时间戳

首页 > 解决方案 > 如何向量化以下python代码

numpy - 如何向量化以下python代码

问题描述

我正在尝试获取一个矩阵，其中每个元素的计算方式如下：

X = torch.ones(batch_size, dim)
X_ = torch.ones(batch_size, dim)
Y = torch.ones(batch_size, dim)
M = torch.zeros(batch_size, batch_size)
for i in range(batch_size):
    for j in range(batch_size):
        M[i, j] = ((X[i] - X_[i] * Y[j])**2).sum()

逐元素计算非常慢M，有没有关于如何使用矩阵乘法来替换 for 循环的建议？

谢谢。

标签： numpymatrixvectorizationmatrix-multiplicationpytorch

解决方案

如果你想sum()过度调暗，你可以将你的 2D 问题“提升”到 3D 并在那里求和：

M = ((X[:, None, :] - X_[:, None, :] * Y[None, ...])**2).sum(dim=2)

它是如何工作的：

X[:, None, :]并且X_[:, None, :]是 3D 的大小(batch_size, 1, dim)，并且Y[None, ...]是大小(1, batch_size, dim)。

当乘以X_[:, None, :] * Y[None, ...]pytorch时，将 size 1 的维度广播到适当的维度以获得 size 的结果(batch_size, batch_size, dim)。
最后，您sum()只需在最后一个维度上获得size(dim=2)的输出。M(batch_size, batch_size)

这里的诀窍是利用广播来完成的。

推荐阅读