Pytorch tensor处理函数/技巧汇总 Cook Blog（持续更新）

二维tensor一次多行/多列 提取行列交叉的技巧（可扩展到三维）

import torch

a = torch.tensor ( [[1,2,3,4],
                    [2,3,1,5],
                    [5,1,7,2]])
print(a[[1,0]])        #一次多行提取
print(a[...,[1,0]])    #一次多列提取
# a[...,[1,0]]等价于a[:,0:2]
print(a[[1,0],[1,2]])  #获得行列交叉的元素

tensor([[2, 3, 1, 5],
        [1, 2, 3, 4]])
tensor([[2, 1],
        [3, 2],
        [1, 5]])
tensor([3, 3])

三维tensor中每一个索引的原点及三轴的方向。

如果将第二个二维tensor放在第一个二维tensor后面的，顶点为正视情况下的左上角，则第一个维度x轴垂直纸面向里（代表厚度），第二个维度由顶点向下（代表行数），第三个维度向右（代表列数）

import torch

a =torch.tensor([[[1,8,0],
                  [3,4,1]],
                 [[2,3,0],
                  [4,5,1]]])
print(a[0][0][0])
print(a[1][0][0])
print(a[0][1][0])
print(a[0][0][1])

tensor(1)
tensor(2)
tensor(3)
tensor(8)

三维tensor沿某一方向进行连乘(torch.prod()函数）

prod中的数字代表着维度，与索引维度一致。如第三维向右，表明三维tensor按照向右的方向进行累乘，最终得到一个2*2 tensor。

import torch

a =torch.tensor([[[1,8,0],
                  [3,4,1]],
                 [[2,3,0],
                  [4,5,1]]])

print(a.prod(0))
print(a.prod(1))
print(a.prod(2))

tensor([[ 2, 24,  0],
        [12, 20,  1]])
tensor([[ 3, 32,  0],
        [ 8, 15,  0]])
tensor([[ 0, 12],
        [ 0, 20]])

三维tensor的两个维度进行交换（torch.transpose()/torch.permute()函数)

使用torch.transpose(dim1,dim2) 函数 代表三维tensor的维度dim1和维度dim2进行交换，从几何角度看是对立方体进行旋转/观察的面发生变化。

torch.transpose() 能够更高效得实现两个维度之间的交换。在数据处理流程中，我们常希望能够对张量的多个维度进行交换，这时需要torch.permute(dim1,dim2,…) 函数。该函数使用示例如下所示：