PyTorch：Tensor的常用基础函数

一、张量维度重构函数view()

函数原型

import torch
Tensor.view(*shape) → Tensor

返回一个新的张量，它的数据与原张量相同但维度不同。

举例

import torch
a = torch.arange(1, 7)
print(a)

b = a.view(2, 3)
print(b)

c = a.view(3, -1)
print(c)

输出为

原张量是1×6的张量，张量b是将a重构成2×3的形状。

view()函数中也可以传入-1作为参数。Tensor.view(-1)会让原张量直接展开成一维结构。如果Torch.view(a, b)中，a和b中有一个传入-1，另一个是正常的参数，说明重构张量的其中一个维度已知，另一个（即-1表示的维度）自动计算出来。

如上例中a.view(3, -1)就是已知第一个维度为3，则自动计算出第二个维度为2，得到c的大小为3×2。

二、高级选择函数

1. 切片函数index_select()

index_select()函数的原型如下。

import torch
torch.index_select(input, dim, index, out=None) → Tensor
torch.Tensor.index_select(dim, index) → Tensor

index_select()函数的作用是针对张量input，在它的dim维度上切取index指定的范围切片。

参数：
input：被操作的张量
dim：维度
index：一维Tensor，表示索引下标的范围

该函数有两种用法，除了将被切片的张量作为参数传入到函数中以外，还可以直接调用张量的内置函数。
例如

import torch
a = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])

# 写法1
b = torch.index_select(a, 0, torch.tensor([1]))
print(b)

# 写法2
c = a.index_select(1, torch.tensor([1,3])) # 等价于c = torch.index_select(a, 1, torch.tensor([1,3]))
print(c)

输出为

这里维度dim从0开始算，则b表示在第0维（即行）上，切下下标为1的行；c表示在第1维（即列）上，切下下标为1和3的列。

2. 批量指定索引函数gather()

函数原型

import torch
torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
torch.Tensor.gather(dim, index) → Tensor

根据indexdim两个参数共同指定的索引，从矩阵input中提取出相应的子阵。

参数：
input：Tensor类型，表示作为来源的张量
dim：int类型，表示索引的维度
index：Tensor类型，表示被聚集的元素的索引

对于一个三维张量，应用gather函数的结果为

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

下面举几个例子来说明这个函数。

import torch
input = torch.arange(4, 13).view(3, 3)
print(input)

index1 = torch.tensor([[2, 1, 0]])
a = torch.gather(input, 0, index1)
print(a)

结果如下

计算方法如下。

1. 首先看index1张量，它是一个大小为1×3的张量，三个元素的位置分别为[0][0]、[0][1]、[0][2]。
2. 然后看dim值为0，即把第0维的索引值替换成对应位置index1张量的值，index1 = [[2, 1, 0]]，因此替换成[2][0]、[1][1]、[0][2]
3. 根据上面的索引值找到input张量中对应元素，取出index1的顺序排列，最后得到的输出张量的size一定与index1相同。即input[2][0]、input[1][1]、input[1][1]，得到的输出为[[10, 8, 6]]
   
   再举一例。

index2 = torch.tensor([[0, 2],[1, 2]])
b = input.gather(1, index2)
print(b)

输出为

过程为

三、数学统计函数

1. 求和函数sum()

求和函数sum()
函数原型：

import torch
torch.sum(input, dim, keepdim=False, *, dtype=None) → Tensor
torch.Tensor.sum(dim=None, keepdim=False, dtype=None) → Tensor 

该函数的作用是返回对输入的张量input的所有元素求和。各参数的含义如下

参数：
input：被操作的张量
dim：int类型，计算进行的维度。
keepdim：bool类型，输出的参数是否继续保持与input张量一致的维度。
dtype：指定返回的张量的类型。默认为None。

举个例子

import torch
a = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])

# 写法1
b1 = torch.sum(a, dim=0)
b2 = torch.sum(a, dim=0, keepdim=True)
print(b1)
print(b2)

# 写法2
c1 = a.sum(dim=1) # 等价于c = torch.sum(a, dim=1)
c2 = a.sum(dim=1, keepdim=True) # 等价于c = torch.sum(a, dim=1, keepdim=True)
print(c1)
print(c2)

d = torch.sum(a)
print(d)

输出如下。

2. 最值函数max()/min()、求最值元素的索引argmax()/argmin()

返回输入的input张量中最大/最小的元素的数值。

import torch
torch.max(input, dim, keepdim=False, *, out=None) → Tensor
torch.Tensor.max(dim=None, keepdim=False) → Tensor

torch.min(input, dim, keepdim=False, *, out=None) → Tensor
torch.Tensor.min(dim=None, keepdim=False) → Tensor

返回输入的input张量中最大/最小的元素的索引值。

import torch
torch.argmax(input, dim, keepdim=False,) → LongTensor
torch.Tensor.argmax(input, dim, keepdim=False,) → LongTensor

torch.argmin(input, dim, keepdim=False,) → LongTensor
torch.Tensor.argmin(dim=None, keepdim=False,) → LongTensor

上述四个函数的用法和参数的含义与sum()函数相同。

四、元素提取函数item()

如果某个张量中只包含一个元素，则可以使用item()函数提取这个元素值。

import torch
a = torch.tensor([2.0])
b = a.item()

print(a)
print(b)

结果如下