时间戳

首页 > 解决方案 > 通过 Trace 对 Pytorch 张量进行排序

sorting - 通过 Trace 对 Pytorch 张量进行排序

问题描述

我有一个 (100,64,22,3,3) 形状的 pytorch 张量，我想通过 (3,3) 分量的轨迹沿轴 = 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]

标签： sortingpytorchvectorizationtensor

解决方案

尝试：

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, :]]

推荐阅读