如何向量化以下python代码

【问题标题】：How to vectorize the following python code如何向量化以下python代码
【发布时间】：2019-04-25 19:02:40
【问题描述】：

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

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 循环的建议？

谢谢。

【问题讨论】：

很抱歉给您带来了困惑。我在这里转移代码时犯了一个错误。实际上 X 是一批向量，所以它的大小是 (batch_size, dim)。所以应用 sum() 将给出一个标量。所以我想我们需要另一个版本的矢量化？

标签： numpy matrix vectorization matrix-multiplication pytorch

【解决方案1】：

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

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

工作原理：

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

当将X_[:, None, :] * Y[None, ...] pytorch broadcasts 大小为1 的维度乘以适当的维度以获得大小为(batch_size, batch_size, dim) 的结果。
最后，您 sum() 仅在最后一个维度 (dim=2) 上得到输出 M 大小为 (batch_size, batch_size)。

这里的技巧是通过利用broadcasting来完成的。

【讨论】：

谢谢！因为 M[i,j] 是一个标量值
感谢您解释代码，它完美地解决了我的问题。