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.transpose(dim1,dim2) 函数 代表三维tensor的维度dim1和维度dim2进行交换，从几何角度看是对立方体进行旋转/观察的面发生变化。代码示例如下：

import torch

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

b = a.transpose(0,1)
d = a.transpose(0,2)
c = a.transpose(1,2)
print(b)
print("  ")
print(c)
print("  ")
print(d)

tensor([[[1, 8, 0],
         [2, 3, 0]],

        [[3, 4, 1],
         [4, 5, 1]]])  
tensor([[[1, 3],
         [8, 4],
         [0, 1]],

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

        [[8, 3],
         [4, 5]],

        [[0, 0],
         [1, 1]]])

二维tensor拼接为三维tensor（torch.stack()函数）低维堆叠到高维

使用torch.stack((tensor1,tensor2,…),dim) 函数实现若干个低维tensor按照高一维之后的dim进行堆叠，这里dim定义的方式同上述三维方向一致。两个tensor有先后顺序之分。对于二维的理解方式为：如dim=2是代表由右上顶点水平向右的方向，则二维tensor1需要与这个方向垂直，同时保持该tensor顶点不变，因此原二维tensor的第一行会变为第一列，第二行会变为垂直纸面向里的第二列。代码示例如下：

import torch

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

print(torch.stack((a,b),dim=0))
print("   ")
print(torch.stack((b,a),dim=0))
print("   ")
print(torch.stack((a,b),dim=1))
print("   ")
print(torch.stack((a,b),dim=2))
tensor([[[1, 8, 0],
         [3, 4, 1]],

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

        [[1, 8, 0],
         [3, 4, 1]]])
   
tensor([[[1, 8, 0],
         [2, 3, 0]],

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

        [[3, 4],
         [4, 5],
         [1, 1]]])

tensor同维度拼接(torch.cat())函数

使用torch.cat((tensor1,tensor2,…),dim) 函数实现将若干个tensor在不改变维度的情况下，同维度之内进行拼接，dim同样指代具体哪个维度进行拼接。示例代码如下：

import torch

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

实现tensor的对位相乘、矩阵乘和有广播机制的矩阵乘法

使用torch.mul(a,b) 实现矩阵a和矩阵b的对位相乘，要求两个矩阵各个维度元素个数完全对应。
使用torch.mm(a,b) 实现矩阵乘法，a、b维度要求与矩阵乘一致。
使用torch.matmul(a,b) 实现带有广播机制的乘法，低维tensor可以依据矩阵乘规则自动适配到高维。

import torch

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

d = torch.tensor([1,8,0])
                  
a_T = a.transpose(0,1) #将a进行转置

print(torch.mul(a,b))
print(torch.mm(a,a_T))
print(torch.matmul(a,d)) #广播操作 将[3] 适配到高维变成了3*1
print(a.size(),d.size()) 

tensor([[ 2, 24,  0],
        [12, 20,  1]])
tensor([[65, 35],
        [35, 26]])
tensor([65, 35])
torch.Size([2, 3]) torch.Size([3])

返回tensor的非0部分(torch.nonzero()函数)

使用torch.nonzero(input,out=None,as_tuple=False) 函数，能够返回input(tensor)中非零的部分的索引，并能将提取input中非零部分的元素。代码示例如下。

import torch

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

d = torch.tensor([1,8,0])
                  
a_T = a.transpose(0,1) #将a进行转置

e = torch.nonzero(a,out=None,as_tuple = False)
e_t = torch.nonzero(a,out=None,as_tuple = True)

print(e)
print(e_t)
print(a[e_t])

tensor([[0, 0],
        [0, 1],
        [1, 0],
        [1, 1],
        [1, 2]])
(tensor([0, 0, 1, 1, 1]), tensor([0, 1, 0, 1, 2]))
tensor([1, 8, 3, 4, 1])

对tensor进行升维/降维(torch.squeeze()/torch.unsqueeze()函数)

在对张量进行处理的过程之中，很多时候涉及到维度的扩增和去除，使用torch.squeeze()/torch.unsqueeze() 函数能够实现对张量维度的增减，squeeze函数用去去除维数为1的维度，unsqueeze用于添加维数为1的维度。代码示例如下

import torch

a = torch.tensor([1,2,3])     # a是一维的tensor
b = torch.tensor([[1,2,3]])   # b是二维的tensor
print(a.shape,b.shape)# a: torch.Size([3]) b:torch.Size([1, 3])

b_2 = a.squeeze(0)
print(b.shape,b_2.shape)
a_2 = a.unsqueeze(0)
print(a.shape,a_2.shape)

torch.Size([1, 3]) torch.Size([3])
torch.Size([3]) torch.Size([1, 3])

两个函数中输入的参数分别代表 删除/增加 第几个维度 0代表第一维度

便捷直接返回tensor的维度(torch.ndimension())函数

直接返回一个tensor的维度，如1,2,3等。其代码示例及等价表示如下所示：

import torch

a = torch.tensor([[1,2,3]])

print(a.ndimension())
print(len(a.shape))

2
2

便生成任意维度的tensor(torch.rand/randn()函数)

便捷生成tensor的pytorch内置函数有很多，之后会逐渐补充。较常用的生成任意维度的函数为torch.rand/randn() 括号内为一组数字，数字代表各个维度的具体数据数目。使用示例如下所示：

import torch

a = torch.rand(2,3,4,5) # 随机生成 具有均值为0 方差为1的正态分布
# b = torch.arange(n)   # 从0~n-1生成一个tensor 如 tensor([0,1,2,3,4,5])
print(a.shape)

torch.Size([2, 3, 4, 5])

对多维度张量按指定维度求均值和方差(torch.mean()、torch.var()函数)

torch在处理均值和方差的内置函数风格上面同numpy非常接近。函数用法分别为torch.mean(tensor,dim,keepdims=True)、torch.var(tensor,dim,keepdims=True)。对四维张量代表的批图像的batchnormalization如下例所示：

import torch

a = torch.randn(2,3,4,5)   #代表一个批处理图像数据 [B,C,H,W]
a_mean = torch.mean(a,(0,2,3,),keepdims=True)  #这里的keepdims代表处理后的数据和原始数据维度是否一致
a_var = torch.var(a,(0,2,3),keepdims=True)

BN = (a - a_mean)/torch.sqrt(a_var)  # 这里在进行计算的时候使用到了tensor的广播机制，对0,2,3通道进行广播
print(a_mean.shape)
print(a_var.shape)
print(BN.shape)

torch.Size([1, 3, 1, 1])
torch.Size([1, 3, 1, 1])
torch.Size([2, 3, 4, 5])

以torch.mean()和torch.var()函数实现的BN简单例子(因为省略了后续的线性变换)就完成啦。

对低维张量进行高维重复扩充(torch.repeat()函数)

在实际定义张量时，我们希望对低维张量在维度不同维度上进行一定的重复，使用torch.repeat(torch.Size）函数可以实现该功能。被重复的对象可以为一维、二维、甚至三维向量。函数的输入分别为各个维度对应重复的次数，对二维分别为行和列，对三维如本章最开始阐述的定义(第一个维度为垂直纸面向里)使用示例如下：

import torch

a = torch.tensor([2,3,4]) # 这是一个一维tensor进行重复的例子
b = a.repeat(4,2,1)
c = a.repeat(4,2)
print(b.shape)
print(b)
torch.Size([4, 2, 3])
tensor([[[2, 3, 4],
         [2, 3, 4]],

        [[2, 3, 4],
         [2, 3, 4]],

        [[2, 3, 4],
         [2, 3, 4]],

        [[2, 3, 4],
         [2, 3, 4]]])
         
a = torch.tensor([[1,2,3],[4,5,6]]) # 这是一个二维tensor进行重复的例子
b = a.repeat(2,1) # 这里如果只想第一个维度重复若干次 则第二维度需为1 该函数输入维度不得小于被处理的张量的维度
print(b)
tensor([[1, 2, 3],
        [4, 5, 6],
        [1, 2, 3],
        [4, 5, 6]])

生成同已知张量维度一致的全零张量(torch.zeros_like())函数

有时为了便于计算，需要生成和已知张量维度一致，所有元素均为0的张量，此时需要用到**torch.zeros_like(tensor1)**函数，返回的张量维度同tensor1一致，元素均为0。示例代码如下：

import torch

a = torch.randn(2,3)
print(a)
b = torch.zeros_like(a)
print(b)

tensor([[-0.9361,  0.7650,  0.3454],
        [ 0.1841,  1.2744,  2.0065]])
tensor([[0., 0., 0.],
        [0., 0., 0.]])

神经网络将张量剥离计算图不进行梯度回传(torch.detach()函数)

有时需要调取神经网络某一层的参数进行一些额外的计算或处理等(如求解loss等)，这时需要将相应的tensor从计算图中剥离出来，不会因为后续的计算进行梯度更新，此处使用**torch.detach()**函数。示例代码如下所示：

import torch

a = torch.tensor([1, 2, 3.], requires_grad=True)
print(a)
b = a.detach()
print(b.requires_grad)

tensor([1., 2., 3.], requires_grad=True)
False