按迹对 Pytorch 张量进行排序

Question

我有一个 (100,64,22,3,3) 形的 pytorch 张量，我想通过 (3,3) 分量的轨迹沿 axis=0 排序。我下面的代码有效，但由于 for 循环，它非常慢。有没有办法矢量化操作以加快速度？

x=torch.rand(100,64,22,3,3)
x_sorted=torch.zeros((x.shape[0],x.shape[1],x.shape[2],x.shape[3],x.shape[4]))
            for i in range(x.shape[0]):
              #compute tensorized trace
              trace=new=torch.diagonal(x[i], dim1=-2, dim2=-1).sum(-1) 
              #Sort the trace
              trace_values,trace_ind=torch.sort(trace,dim=0,descending=True)
              for j in range(x_sorted.shape[1]):
                for k in range(x_sorted.shape[2]):
                  x_sorted[i,j,k]=x[i,trace_ind[j,k],k]

Answer 1

尝试：

tensor = torch.tensor(np.random.rand(100,64, 3, 3))

orders = torch.argsort(torch.einsum('ijkk->ijk', tensor).sum(-1), axis=0)
orders.shape

tensor[orders, torch.arange(s.shape[1])[None, :]]

按迹对 Pytorch 张量进行排序

Sorting a Pytorch Tensor by Trace

sorting

vectorization

pytorch

tensor