tensor笔记1--index_select函数(对tensor维度的理解)

pytorch中index_select函数的作用选取某一维度上的数据

函数形式为:

index_select(input, dim, index)

input为tensor,dim是维度从0开始,index是一维tensor(向量),表示在这个维度上要选择的下标

下边直接上例子:

import torch

x = torch.Tensor([[[1, 2, 3],
                   [4, 5, 6]],

                  [[9, 8, 7],
                   [6, 5, 4]]])
print(x)
print(x.size())
index = torch.LongTensor([0, 0, 1])
print(torch.index_select(x, 0, index))
print(torch.index_select(x, 0, index).size())
print(torch.index_select(x, 1, index))
print(torch.index_select(x, 1, index).size())
print(torch.index_select(x, 2, index))
print(torch.index_select(x, 2, index).size())

input的张量形状为2×2×3,index为[0, 0, 1]的向量

分别从0、1、2三个维度来使用index_select()函数,并输出结果和形状,维度大于2就会报错因为input最大只有三个维度

输出:

tensor([[[1., 2., 3.],
         [4., 5., 6.]],

        [[9., 8., 7.],
         [6., 5., 4.]]])
torch.Size([2, 2, 3])
tensor([[[1., 2., 3.],
         [4., 5., 6.]],

        [[1., 2., 3.],
         [4., 5., 6.]],

        [[9., 8., 7.],
         [6., 5., 4.]]])
torch.Size([3, 2, 3])
tensor([[[1., 2., 3.],
         [1., 2., 3.],
         [4., 5., 6.]],

        [[9., 8., 7.],
         [9., 8., 7.],
         [6., 5., 4.]]])
torch.Size([2, 3, 3])
tensor([[[1., 1., 2.],
         [4., 4., 5.]],

        [[9., 9., 8.],
         [6., 6., 5.]]])
torch.Size([2, 2, 3])

对结果进行分析:

index是大小为3的向量,输入的张量形状为2×2×3

dim = 0时,输出的张量形状为3×2×3

dim = 1时,输出的张量形状为2×3×3

dim = 2时,输出的张量形状为2×2×3

注意输出张量维度的变化与index大小的关系,结合输出的张量与原始张量来分析index_select()函数的作用

 

你可能感兴趣的:(python,机器学习)