在pytorch中有效地求和向量对

Question

我正在尝试计算矩阵中每对行的总和。假设我有一个 m x 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 循环的情况下为任意矩阵创建这样的矩阵？

Answer 1

你可以试试这个功能：

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 循环，AFAIK，但应该为可变大小的输入完成这项工作。