首先要感谢CSDN中http://t.csdn.cn/XyT4e这篇文章(我接下来写的内容,也和这篇文章基本一样)
下面是我实际操作得到的结果:
我们看第一种情况的代码:
import torch b = torch.arange(1, 61).reshape(3, 4, 5) idx1 = torch.tensor([0, 0, 2]).unsqueeze(-1).repeat(1, 4) bb = b[idx1, : print(bb)
我们先来看一下张量b的内容:
''' 张量b内容: tensor([[[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[21, 22, 23, 24, 25], [26, 27, 28, 29, 30], [31, 32, 33, 34, 35], [36, 37, 38, 39, 40]], [[41, 42, 43, 44, 45], [46, 47, 48, 49, 50], [51, 52, 53, 54, 55], [56, 57, 58, 59, 60]]])) '''
接下来我们看一下索引得到的张量bb的内容(这个有点长):
''' 张量bb的内容为: tensor([[[[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]], [[[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]], [[ 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]], [[[41, 42, 43, 44, 45], [46, 47, 48, 49, 50], [51, 52, 53, 54, 55], [56, 57, 58, 59, 60]], [[41, 42, 43, 44, 45], [46, 47, 48, 49, 50], [51, 52, 53, 54, 55], [56, 57, 58, 59, 60]], [[41, 42, 43, 44, 45], [46, 47, 48, 49, 50], [51, 52, 53, 54, 55], [56, 57, 58, 59, 60]], [[41, 42, 43, 44, 45], [46, 47, 48, 49, 50], [51, 52, 53, 54, 55], [56, 57, 58, 59, 60]]]]) '''
首先b是(3,4,5)的张量,然后b[idx1, :, ]这里索引就是将张量b的后面两个维度(4, 5)当作一个整体。然后根据idx1中的内容进行索引。又因为idx1的shape是(3, 4)
所以索引后的bb的shape为(3, 4, 4, 5)。然后也就得到了那个结果。
第二种情况的代码为(这里我们的张量b仍选用第一种情况的张量b):
import torch b = torch.arange(1, 61).reshape(3, 4, 5) idx1 = torch.tensor([0, 0, 2]).unsqueeze(-1).repeat(1, 4) idx2 = torch.randint(0, 4, (3, 4), dtype = torch.long) cc = b[idx1, idx2] cc.shape,cc
此时的输出结果为:
''' idx2得到的随机值为: tensor([[1, 2, 1, 3], [2, 1, 1, 2], [2, 1, 0, 3]]) cc.shape: torch.Size([3, 4, 5]) 张量cc为: tensor([[[ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [ 6, 7, 8, 9, 10], [16, 17, 18, 19, 20]], [[11, 12, 13, 14, 15], [ 6, 7, 8, 9, 10], [ 6, 7, 8, 9, 10], [11, 12, 13, 14, 15]], [[51, 52, 53, 54, 55], [46, 47, 48, 49, 50], [41, 42, 43, 44, 45], [56, 57, 58, 59, 60]]]) '''
此时我们可以看到idx1与idx2他们都是(3,4)d的矩阵,所以对于b[idx1, idx2]会先将idx1与idx2组合起来然后作为索引,去索引张量b中的第三个维度[5]。然后就得到结果了。
1.如果我们要多维度的索引，我们需要保证dim(idx1) = dim(idx2)
2.多个维度一起索引时,我们先把两个维度叠加在一起,然后根据新的索引去索引