pytorch小知识点（三）-------Tensor的indices操作

短命女 2022-02-12 19:15 395阅读 0赞

# 一、Tesnor #

先简单介绍一下Tensor。Tensor是pytorch的核心，它是一个包含单一数据类型的多维矩阵。

pyTorch定义了七种CPU tensor类型和八种GPU tensor类型：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2dvb2R4aW5faWU_size_16_color_FFFFFF_t_70][]

默认的torch.Tensor是FloatTensor。我们可以简单地向下面的方式创建一个Tensor:

"""
    FloatTensor
    """
    x1 = torch.FloatTensor([1,2,3,4,5,6])
    
    x1
    Out[100]: tensor([1., 2., 3., 4., 5., 6.])
    
    x1.dtype
    Out[115]: torch.float32
    
    
    
    """
    LongTensor
    """
    x2 = torch.LongTensor([1,2,3,4,5,6])
    
    x2
    Out[116]: tensor([1, 2, 3, 4, 5, 6])
    
    x2.dtype
    Out[117]: torch.int64

# 二、Tensor的indices操作 #

我们都知道对于Tesnor可以向numpy一样进行slice（切片）操作，如下：

x = torch.randn((3,4))
    
    x
    Out[119]: 
    tensor([[-0.2558,  0.9740, -0.4165,  0.0370],
            [-0.6938, -0.6043, -0.1243, -1.0082],
            [ 0.4293,  1.1933,  0.6975,  0.1752]])
    
    x[1,:]       #取第二行
    Out[120]: tensor([-0.6938, -0.6043, -0.1243, -1.0082])
    
    x[:,2]       #取第三列
    Out[121]: tensor([-0.4165, -0.1243,  0.6975])

但是Tensor中有一个独特的indices操作，按照我们给的索引进行取数：

形式为x\[mask\]

*  x要求为1个Tensor
 *  mask必须为ByteTensor或者LongTensor(注意：这种格式的mask对应的结果是完全不同的)

### 1.当mask为ByteTensor时，要求**mask必须和x形状相同**。 ###

此时mask表示是否保留x中对应位置的数。若mask中对应位不为0，则保留x中对应数。最终返回值是一个1-D Tensor，长度等于mask中非0元素的个数(因此，建议大家在mask中只用0和1两种数字就可以了)

e.g

x = torch.Tensor([[0,1,2,3,4],[5,6,7,8,9],[10,11,12,13,14]])
    
    
    """
    mask1 对应保留x中的奇数
    """
    mask1 = torch.ByteTensor([[0,1,0,1,0],[0,1,0,1,0],[0,1,0,1,0]])
    
    """
    mask2  对应保留x中的偶数
    """
    mask2 = torch.ByteTensor([[1,0,1,0,1],[1,0,1,0,1],[1,0,1,0,1]])
    
    x[mask1]
    Out[148]: tensor([ 1.,  3.,  6.,  8., 11., 13.])
    
    x[mask2]
    Out[149]: tensor([ 0.,  2.,  4.,  5.,  7.,  9., 10., 12., 14.])
    
    
    """
    mask3  只是为了证明mask中所有非0数均起相同作用，此时mask1和mask3的实质作用相同
    但是不推荐使用mask3这种存在0和1之外其他数的形式，mask1表达的意义明显比mask3清晰
    """
    
    mask3 = torch.ByteTensor([[0,1,0,3,0],[0,6,0,8,0],[0,11,0,13,0]])
    
    x[mask3]
    Out[151]: tensor([ 1.,  3.,  6.,  8., 11., 13.])

### 2.当mask为LongTensor时，此时x和mask形状不必相同。 ###

注：当mask为LongTensor时，如果不熟悉的话最好使用1-D Tensor的mask

我们先看看x为1-D，mask为1-D Tensor的情况

x1 = torch.Tensor([0,1,2,3,4])
    mask1 = torch.LongTensor([1,0,2])
    
    x1[mask1]
    
    Out[174]: tensor([1., 0., 2.])
    """
    为什么输出会是[1,0,2]？？？
    """
    
    mask2 = torch.LongTensor([4,3,2,1,0])
    
    x1[mask2]
    
    Out[176]: tensor([4., 3., 2., 1., 0.])
    
    """
    到这儿我们该看懂了吧，mask为LongTensor时,x[mask]相当于一个在0纬度上的gather操作
    就是按照mask中的索引对应的x中的数字放到索引当前位置，用mask1来举例，mask1中第一个数为1，我们就在结果的第一个数上放x1[1],mask1中第二个数为0,结果的第二个数就是x1[0]
    因此，mask1对应的输出就是[x1[mask1[0],x1[mask1[1],x1[mask1[2]]
    此时不难理解mask2对应的输出为[4., 3., 2., 1., 0.]了吧
    """

那么x为2-D，mask为1-D的情况呢？

x2 = torch.Tensor([[ 0.,  1.,  2.,  3.,  4.],
                       [ 5.,  6.,  7.,  8.,  9.],
                       [10., 11., 12., 13., 14.]])
    
    mask1 = torch.LongTensor([0,1,2])
    
    x2[mask1]
    
    Out[182]: 
    tensor([[ 0.,  1.,  2.,  3.,  4.],
            [ 5.,  6.,  7.,  8.,  9.],
            [10., 11., 12., 13., 14.]])
    
    mask2 = torch.LongTensor([0,2,1])
    
    x2[mask2]
    Out[184]: 
    tensor([[ 0.,  1.,  2.,  3.,  4.],
            [10., 11., 12., 13., 14.],
            [ 5.,  6.,  7.,  8.,  9.]])
    
    """
    此时我们可以得出结论，不论x是几维的，mask为LongTensor时都是按照mask中的顺序排列x[mask[i]]
    上面的mask1对应的结果应该是 [x2[0],x2[1],x2[2]]
    mask2对应的结果是[x2[0],x2[1],x2[2]]
    直白来说这种操作就是对x中的第一个维度上的重排列
    """

有了上面的铺垫，我们可以看一下当mask为多维时的情况

x2 = torch.Tensor([[ 0.,  1.,  2.,  3.,  4.],
                       [ 5.,  6.,  7.,  8.,  9.],
                       [10., 11., 12., 13., 14.]])
    
    mask1 = torch.LongTensor([[0,1,2],[0,1,2]])
    
    x2[mask1]
    
    Out[188]: 
    tensor([[[ 0.,  1.,  2.,  3.,  4.],
             [ 5.,  6.,  7.,  8.,  9.],
             [10., 11., 12., 13., 14.]],
    
            [[ 0.,  1.,  2.,  3.,  4.],
             [ 5.,  6.,  7.,  8.,  9.],
             [10., 11., 12., 13., 14.]]])
    
    
    mask2 = torch.LongTensor([[0,0,0],[0,0,0]])
    
    x2[mask2]
    Out[190]: 
    tensor([[[0., 1., 2., 3., 4.],
             [0., 1., 2., 3., 4.],
             [0., 1., 2., 3., 4.]],
    
            [[0., 1., 2., 3., 4.],
             [0., 1., 2., 3., 4.],
             [0., 1., 2., 3., 4.]]])
    
    """
    what?什么情况，输出怎么变得这么复杂？冷静下来仔细分析一下，
    x2[mask1]看起来是将x2在一个新的维度上扩展了两倍 ，x2[mask2]好像是将x2[0]先复制三遍，再在一个新的维度上扩展了两倍
    貌似有点规律。我们还记得上上面的结论吗，mask为LongTensor时就是将x[mask[i]]按照mask中的顺序排列
    我们试着将mask1和mask2中的数字i替换成x2[i]看看是不是就是输出结果呢？
    其实，不论mask多么复杂，我们将mask中的数字i当作x在第一个纬度上的索引就好了，然后将i在原位置替换成x[i]就是最终的输出结果了
    """
    
    """
    我们来测试一下上面的结论
    x 为1-D mask为 2-D
    """
    
    x3 = torch.Tensor([0.1,0.2,0.3])
    
    mask1 = torch.LongTensor([[0,1,2],[0,1,2]])
    
    x3[mask1]
    Out[193]: 
    tensor([[0.1000, 0.2000, 0.3000],
            [0.1000, 0.2000, 0.3000]])
    
    mask2 = torch.LongTensor([[0,0,0],[0,0,0]])
    
    x3[mask2]
    Out[194]: 
    tensor([[0.1000, 0.1000, 0.1000],
            [0.1000, 0.1000, 0.1000]])

--------------------

x\[mask\]中mask为LongTensor时与gather操作的不同

x\[mask\]中的mask始终指代的是x的第一个维度上的索引。

gather中的index可以为指定dim上的索引，而且对于gather操作来说，index必须和操作数shape相同

--------------------

# 三、indices为ByteTensor的应用实例 #

这个实例完成的功能是在一个2-D Tensor中截取每一行最大的前ni(每一行的nibuxiangdeng)个数

我们有一个4\*5的Tensor

x = torch.randn((4,5))
    
    x
    Out[199]: 
    tensor([[ 1.8299,  0.0714, -0.1213, -0.5910, -0.2492],
            [-1.5000, -0.5365,  0.1588,  1.1101,  0.8724],
            [ 1.1998,  2.3945,  1.0767,  0.5229,  0.3283],
            [-1.0195,  0.3819, -1.0285,  1.2028, -1.0717]])

我们来截取他每行的前\[1,2,3,4\]个数，即取第一行最大的1个数，第二行2个...

x = torch.randn((4,5))
    """
    x
    Out[199]: 
    tensor([[ 1.8299,  0.0714, -0.1213, -0.5910, -0.2492],
            [-1.5000, -0.5365,  0.1588,  1.1101,  0.8724],
            [ 1.1998,  2.3945,  1.0767,  0.5229,  0.3283],
            [-1.0195,  0.3819, -1.0285,  1.2028, -1.0717]])
    """
    _,indices = x.sort(dim=1,descending=True)
    
    """
    indices
    Out[201]: 
    tensor([[0, 1, 2, 4, 3],
            [3, 4, 2, 1, 0],
            [1, 0, 2, 3, 4],
            [3, 1, 0, 2, 4]])
    """
    _,idx = indices.sort(dim=1)
    """
    idx
    Out[203]: 
    tensor([[0, 1, 2, 4, 3],
            [4, 3, 2, 0, 1],
            [1, 0, 2, 3, 4],
            [2, 1, 3, 0, 4]])
    
    """
    num = torch.LongTensor([[1],[2],[3],[4]])
    """
    num
    Out[205]: 
    tensor([[1],
            [2],
            [3],
            [4]])
    """
    num = num.expand_as(idx)
    
    """
    num
    Out[207]: 
    tensor([[1, 1, 1, 1, 1],
            [2, 2, 2, 2, 2],
            [3, 3, 3, 3, 3],
            [4, 4, 4, 4, 4]])
    """
    mask = idx<num
    """
    mask
    Out[208]: 
    tensor([[1, 0, 0, 0, 0],
            [0, 0, 0, 1, 1],
            [1, 1, 1, 0, 0],
            [1, 1, 1, 1, 0]], dtype=torch.uint8)
    
    x[mask]
    Out[209]: 
    tensor([ 1.8299,  1.1101,  0.8724,  1.1998,  2.3945,  1.0767, -1.0195,  0.3819,
            -1.0285,  1.2028])
    """

上面例子中的前面的两次排序操作都是为了获得最终的mask，最终的mask为1的地方就是x中最大的几个数（每行的前ni个)。

上面例子还有一个重要的知识点，就是对一个**2-D Tensor使用两次sort函数（对原矩阵、第一次sort之后的index）后，index（第二个变量）的输出是原矩阵一个由大到小（或由小到大）的映射，映射的值为整数，位置保持不变**。

我们可以用个例子体会一下：

x
    Out[215]: 
    tensor([[ 0.2849, -1.0989, -0.2956,  1.2607, -1.1120],
            [ 0.0420,  0.3557, -0.7565,  0.5844, -2.0583],
            [ 0.6427, -1.5065,  0.9481, -1.2509,  2.2974],
            [ 0.1249, -1.1839, -0.3964,  0.3081, -0.3041]])
    
    _,indices = x.sort(dim=1)
    
    _,idx = indices.sort(dim=1)
    
    idx
    Out[218]: 
    tensor([[3, 1, 2, 4, 0],     
            [2, 3, 1, 4, 0],
            [2, 0, 3, 1, 4],
            [3, 0, 1, 4, 2]])

参考资料：\[1\][找出矩阵中升序或降序元素的位置][Link 1]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2dvb2R4aW5faWU_size_16_color_FFFFFF_t_70]: /images/20220212/1db701b3b9794d039ea6d64518ca7a90.png
[Link 1]: https://blog.csdn.net/LXX516/article/details/78804884