在pytorch中有效地求和向量对



我试图计算矩阵中每对行的总和。假设我有一个m × n矩阵,比如

[[1,2,3],
[4,5,6],
[7,8,9]]

,我想创建一个矩阵所有行对的和。因此,对于上面的矩阵,我们想要

[[5,7,9],
[8,10,12],
[11,13,15]]

一般来说,我认为新的矩阵将是(m选择2)x n。对于上面的pytorch示例,我运行

import torch
x = torch.tensor([[1,2,3], [4,5,6], [7,8,9]])
y = x[None] + x[:, None]
torch.cat((y[0, 1:3, :], y[1, 2:3, :]))

手动创建我正在寻找的矩阵。然而,我正在努力思考一种不手动指定索引和不使用for循环来创建输出的方法。有没有一种方法可以在不使用for循环的情况下为任意矩阵创建这样一个矩阵?

您可以尝试使用这个函数:

def sum_rows(x): 
y = x[None] + x[:, None] 
ind = torch.tril_indices(x.shape[0], x.shape[0], offset=-1)
return y[ind[0], ind[1]]

因为你知道你想要具有sum_matrix[i,j], where i<j约束的对(但i>j也可以),你可以指定你想要3D矩阵的下/上三角形指标。这仍然使用for循环,但应该可以处理可变大小的输入。

最新更新