2-张量的API上

1. torch.chunk

尝试将一个张量分割成指定数量的块。每个块都是输入张量的一个视图。

# 创建一个5行6列的张量
x_input_chunk = torch.arange(30).reshape(5,6)
# 将x_input_chunk按行分割成3分，如果除不尽，多余的在最后一个
x_chunk_dim_0 = x_input_chunk.chunk(3,dim=0)
# 将x_input_chunk按列分割成3分，如果除不尽，多余的在最后一个
x_chunk_dim_1 = x_input_chunk.chunk(3,dim=1)
print(f"x_input_chunk={x_input_chunk}")
print(f"x_chunk_dim_0={x_chunk_dim_0}")
print(f"x_chunk_dim_1={x_chunk_dim_1}")

x_input_chunk=tensor([[ 0,  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]])

# 将x_input_chunk按行分割成3分，如果除不尽，多余的在最后一个
x_chunk_dim_0=(tensor([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11]]), tensor([[12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23]]), tensor([[24, 25, 26, 27, 28, 29]]))
# 将x_input_chunk按列分割成3分，如果除不尽，多余的在最后一个
x_chunk_dim_1=(tensor([[ 0,  1],
        [ 6,  7],
        [12, 13],
        [18, 19],
        [24, 25]]), tensor([[ 2,  3],
        [ 8,  9],
        [14, 15],
        [20, 21],
        [26, 27]]), tensor([[ 4,  5],
        [10, 11],
        [16, 17],
        [22, 23],
        [28, 29]]))

2. torch.tensor_split

根据索引或indices_or_sections指定的节数，沿着维dim将张量拆分为多个子张量，所有子张量都是输入视图。将张量按照指定的份额进行分割；
torch.tensor_split(z,(1,6),dim=1)：表示将张量z按照列进行分割，第一刀在序号1上，第2 在序号6上，刀下在序号的左边。

import torch

x = torch.arange(8)
# 将张量X切割成3份
x_tensor_split = torch.tensor_split(x,3)
print(f"x={x}")
print(f"x.tensor_split_3={x_tensor_split}")
y = torch.arange(7)
# 将张量y切割成3份
y_tensor_split_3 = torch.tensor_split(y, 3)
print(f"y={y}")
print(f"y_tensor_split_3={y_tensor_split_3}")
# 将张量x 切割，第一刀在第1位，第二刀在第6位
# 下刀位为下标的左位
# [0,1,2,3,4,5,6,7]->([0],[1,2,3,4,5],[6,7])
x_tensor_split_1_6 = torch.tensor_split(x,(1,6))
print(f"x_tensor_split_1_6={x_tensor_split_1_6}")
z = torch.arange(14).reshape(2,7)
# 将 z 按照列切割成 3 份
z_3_dim_1 = torch.tensor_split(z, 3, dim=1)
# 将 z 按照第一刀在序号1，第二到在序号6的列方向进行分割
z_1_6_dim_1 = torch.tensor_split(z,(1,6),dim=1)
print(f"z={z}")
print(f"z_3_dim_1={z_3_dim_1}")
print(f"z_1_6_dim_1={z_1_6_dim_1}")

x=tensor([0, 1, 2, 3, 4, 5, 6, 7])
x.tensor_split_3=(tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7]))
y=tensor([0, 1, 2, 3, 4, 5, 6])
y_tensor_split_3=(tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6]))
x_tensor_split_1_6=(tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6, 7]))
z=tensor([[ 0,  1,  2,  3,  4,  5,  6],
        [ 7,  8,  9, 10, 11, 12, 13]])
z_3_dim_1=(tensor([[0, 1, 2],
        [7, 8, 9]]), tensor([[ 3,  4],
        [10, 11]]), tensor([[ 5,  6],
        [12, 13]]))
z_1_6_dim_1=(tensor([[0],
        [7]]), tensor([[ 1,  2,  3,  4,  5],
        [ 8,  9, 10, 11, 12]]), tensor([[ 6],
        [13]]))
m=tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])
m_split=(tensor([0, 1, 2]), tensor([3, 4, 5, 6, 7]), tensor([ 8,  9, 10, 11]))

3. torch.dsplit & torch.hsplit

区别
举例，对于t=(2,4,6)来说

• hsplit : 表示按照第dim=1维进行分割，按每个批量中4行6列矩阵中按行进行分割
• dsplit：表示按照dim=2维进行分割，按每个批量中4行6列矩阵中按列进行分割

import torch
n = torch.arange(24.).reshape(2,3,4)
# 按照最后一维进行切割，第一刀在序号为1位置，第二刀在第3位置，刀在序号左侧

n_dplit = torch.dsplit(n,(1,3))
n_t_split = torch.tensor_split(n,(1,3),dim=2)
print(f"n={n}")
print(f"n_dplit={n_dplit}")
print(f"n_t_split={n_t_split}")

n=tensor([[[ 0.,  1.,  2.,  3.],
         [ 4.,  5.,  6.,  7.],
         [ 8.,  9., 10., 11.]],

        [[12., 13., 14., 15.],
         [16., 17., 18., 19.],
         [20., 21., 22., 23.]]])
n_dplit=(tensor([[[ 0.],
         [ 4.],
         [ 8.]],

        [[12.],
         [16.],
         [20.]]]), tensor([[[ 1.,  2.],
         [ 5.,  6.],
         [ 9., 10.]],

        [[13., 14.],
         [17., 18.],
         [21., 22.]]]), tensor([[[ 3.],
         [ 7.],
         [11.]],

        [[15.],
         [19.],
         [23.]]]))
n_t_split=(tensor([[[ 0.],
         [ 4.],
         [ 8.]],

        [[12.],
         [16.],
         [20.]]]), tensor([[[ 1.,  2.],
         [ 5.,  6.],
         [ 9., 10.]],

        [[13., 14.],
         [17., 18.],
         [21., 22.]]]), tensor([[[ 3.],
         [ 7.],
         [11.]],

        [[15.],
         [19.],
         [23.]]]))

t = torch.arange(48.0).reshape(2,4,6)
# hsplit 相当于每个批量的矩阵按照行切割成2份
# t_hsplit[0]=(2,2,6);t_hsplit[1]=(2,2,6)
t_hsplit = torch.hsplit(t, 2)
# dsplit 相当于每个批量的举证按照列切割成2份
# t_dsplit[0]=(2,4,3);t_dsplit[1]=(2,4,3)
t_dsplit = torch.dsplit(t, 2)
print(f"t={t}")
print(f"t_hsplit={t_hsplit}")
print(f"t_hsplit[0]={t_hsplit[0]}")
print(f"t_hsplit[0].shape={t_hsplit[0].shape}")
print(f"t_hsplit[1]={t_hsplit[1]}")
print(f"t_hsplit[1].shape={t_hsplit[1].shape}")
print(f"t_dsplit={t_dsplit}")
print(f"t_dsplit[0]={t_dsplit[0]}")
print(f"t_dsplit[0].shape={t_dsplit[0].shape}")
print(f"t_dsplit[1]={t_dsplit[1]}")
print(f"t_dsplit[1].shape={t_dsplit[1].shape}")

输出结果

t=tensor([[[ 0.,  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.]]])
t_hsplit=(tensor([[[ 0.,  1.,  2.,  3.,  4.,  5.],
         [ 6.,  7.,  8.,  9., 10., 11.]],

        [[24., 25., 26., 27., 28., 29.],
         [30., 31., 32., 33., 34., 35.]]]), tensor([[[12., 13., 14., 15., 16., 17.],
         [18., 19., 20., 21., 22., 23.]],

        [[36., 37., 38., 39., 40., 41.],
         [42., 43., 44., 45., 46., 47.]]]))
t_hsplit[0]=tensor([[[ 0.,  1.,  2.,  3.,  4.,  5.],
         [ 6.,  7.,  8.,  9., 10., 11.]],

        [[24., 25., 26., 27., 28., 29.],
         [30., 31., 32., 33., 34., 35.]]])
t_hsplit[0].shape=torch.Size([2, 2, 6])
t_hsplit[1]=tensor([[[12., 13., 14., 15., 16., 17.],
         [18., 19., 20., 21., 22., 23.]],

        [[36., 37., 38., 39., 40., 41.],
         [42., 43., 44., 45., 46., 47.]]])
t_hsplit[1].shape=torch.Size([2, 2, 6])
t_dsplit=(tensor([[[ 0.,  1.,  2.],
         [ 6.,  7.,  8.],
         [12., 13., 14.],
         [18., 19., 20.]],

        [[24., 25., 26.],
         [30., 31., 32.],
         [36., 37., 38.],
         [42., 43., 44.]]]), tensor([[[ 3.,  4.,  5.],
         [ 9., 10., 11.],
         [15., 16., 17.],
         [21., 22., 23.]],

        [[27., 28., 29.],
         [33., 34., 35.],
         [39., 40., 41.],
         [45., 46., 47.]]]))
t_dsplit[0]=tensor([[[ 0.,  1.,  2.],
         [ 6.,  7.,  8.],
         [12., 13., 14.],
         [18., 19., 20.]],

        [[24., 25., 26.],
         [30., 31., 32.],
         [36., 37., 38.],
         [42., 43., 44.]]])
t_dsplit[0].shape=torch.Size([2, 4, 3])
t_dsplit[1]=tensor([[[ 3.,  4.,  5.],
         [ 9., 10., 11.],
         [15., 16., 17.],
         [21., 22., 23.]],

        [[27., 28., 29.],
         [33., 34., 35.],
         [39., 40., 41.],
         [45., 46., 47.]]])
t_dsplit[1].shape=torch.Size([2, 4, 3])

• hstack：表示按照第dim=1维进行堆砌在一起，相当于每个批次的行堆砌
• dstack：表示按照第dim=2维进行堆砌在一起，相当于每个批次的列堆砌

a = torch.tensor([1,2,3])
b = torch.tensor([4,5,6])
c = torch.hstack((a,b))
print(f"a={a}")
print(f"b={b}")
print(f"c={c}")
a_1 = torch.ones(3,4)
b_1 = torch.zeros(3,4)
c_1 = torch.hstack((a_1,b_1))
print(f"a_1={a_1}")
print(f"b_1={b_1}")
print(f"c_1={c_1}")
print(f"a_1.shape={a_1.shape}")
print(f"b_1.shape={b_1.shape}")
print(f"c_1.shape={c_1.shape}")
a_2 = torch.ones((2,3,4))
b_2 = torch.zeros((2,3,4))
c_2 = torch.hstack((a_2,b_2))
c_3 = torch.dstack((a_2,b_2))
print(f"a_2={a_2}")
print(f"b_2={b_2}")
print(f"c_2_hstack={c_2}")
print(f"c_3_dstack={c_3}")
print(f"a_2.shape={a_2.shape}")
print(f"b_2.shape={b_2.shape}")
print(f"c_2_hstack.shape={c_2.shape}")
print(f"c_3_dstack.shape={c_3.shape}")

a=tensor([1, 2, 3])
b=tensor([4, 5, 6])
c=tensor([1, 2, 3, 4, 5, 6])
a_1=tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
b_1=tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
c_1=tensor([[1., 1., 1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 1., 0., 0., 0., 0.]])
a_1.shape=torch.Size([3, 4])
b_1.shape=torch.Size([3, 4])
c_1.shape=torch.Size([3, 8])
a_2=tensor([[[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]],

        [[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]]])
b_2=tensor([[[0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]],

        [[0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]]])
c_2_hstack=tensor([[[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]],

        [[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]]])
c_3_dstack=tensor([[[1., 1., 1., 1., 0., 0., 0., 0.],
         [1., 1., 1., 1., 0., 0., 0., 0.],
         [1., 1., 1., 1., 0., 0., 0., 0.]],

        [[1., 1., 1., 1., 0., 0., 0., 0.],
         [1., 1., 1., 1., 0., 0., 0., 0.],
         [1., 1., 1., 1., 0., 0., 0., 0.]]])
a_2.shape=torch.Size([2, 3, 4])
b_2.shape=torch.Size([2, 3, 4])
c_2_hstack.shape=torch.Size([2, 6, 4])
c_3_dstack.shape=torch.Size([2, 3, 8])

4. torch.gather

沿dim指定的轴收集值。举例说明

import torch
new_t = torch.tensor([[1, 2], [3, 4]])
output_t_0 = torch.gather(new_t, 0, torch.tensor([[0, 0], [1, 0]]))
output_t_1 = torch.gather(new_t, 1, torch.tensor([[0, 0], [1, 0]]))
print(f"new_t={new_t}")
# new_t=tensor([[1, 2],
#         [3, 4]])
print(f"output_t_0={output_t_0}")
# output_t_0=tensor([[1, 2],
#        [3, 2]])
print(f"output_t_1={output_t_1}")
# output_t_1=tensor([[1, 1],
#        [4, 3]])

当我们输入的矩阵new_t表示如下
$$ne{w}_t=\left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right]$$
$$[0,0]=1;[0,1]=2;[1,0]=3;[1,1]=4$$

• torch.gather(new_t,0,torch.tensor([[0,0],[1,0]]))
  表示的是我们将dim=0替换成给定的坐标，其他的dim保持不变
  $$\left[\begin{array}{cc}0& 0\\ 1& 0\end{array}\right]->\left[\begin{array}{cc}[0,x_1]& [0,x_2]\\ [1,x_3]& [0,x_4]\end{array}\right]->\left[\begin{array}{cc}[0,x_1=0]& [0,x_2=1]\\ [1,x_3=0]& [0,x_4=1]\end{array}\right]->\left[\begin{array}{cc}[0,0]& [0,1]\\ [1,0]& [0,1]\end{array}\right]$$
  我们在将通过索引去在输入张量中找到相关数据
  $$[0,0]=1;[0,1]=2;[1,0]=3,[0,1]=2$$
  所以 output_t_0 = torch.gather(new_t,0,torch.tensor([[0,0],[1,0]]))=torch.tensor([[1,2],[3,2]])
• torch.gather(new_t,1,torch.tensor([[0,0],[1,0]]))
  表示的是我们将dim=1替换成给定的坐标，其他的dim保持不变
  $$\left[\begin{array}{cc}0& 0\\ 1& 0\end{array}\right]->\left[\begin{array}{cc}[y_1,0]& [y_2,0]\\ [y_3,1]& [y_4,0]\end{array}\right]->\left[\begin{array}{cc}[y_1=0,0]& [y_2=0,0]\\ [y_3=1,1]& [y_4=1,0]\end{array}\right]->\left[\begin{array}{cc}[0,0]& [0,0]\\ [1,1]& [1,0]\end{array}\right]$$
  我们在将通过索引去在输入张量中找到相关数据
  $$[0,0]=1;[0,0]=1;[1,1]=4,[1,0]=3$$
  所以 output_t_1= torch.gather(new_t,1,torch.tensor([[0,0],[1,0]]))=torch.tensor([[1,1],[4,3]])
  以上手动计算的跟程序结果一致！

5. torch.scatter_

这是与gather()中描述的方式相反的操作。

Tensor.scatter_(dim, index, src, reduce=None) → Tensor

• dim：索引所沿的轴
• index: 要分散的元素的索引可以是空的，也可以是与src相同维数的。当该操作为空时，该操作将不改变返回值。
• src:要分散的源元素。
• reduce:要应用的减数运算，可以是“加”或“乘”
  举例来说：
  假设我们有一个原始张量src如下：
• 原始张量src
  $$src=torch.tensor([[1,2,3,4,5],[6,7,8,9,10]])$$
• 索引张量index:
  $$index=torch.tensor([[[0,1,2,0]])$$
• 按照dim=0形式将src放到新的全零(3,5)张量中

import torch
src = torch.arange(1, 11).reshape(2, 5)

# src=tensor([[ 1,  2,  3,  4,  5],
#         [ 6,  7,  8,  9, 10]])

index = torch.tensor([[0, 1, 2, 0]])
# index=tensor([[0, 1, 2, 0]])
output = torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src)
# output=tensor([[1, 0, 0, 4, 0],
#         [0, 2, 0, 0, 0],
#         [0, 0, 3, 0, 0]])
print(f"src={src}")
print(f"index={index}")
print(f"output={output}")

其实按照gather方式一样，dim=0表示的第dim=0维度为指定的index的值，其他的保持不变,以前多少就多少列
$$[[0,1,2,0]]->[[(0,0),(1,1),(2,2),(0,3)]]$$
也就是说，需要将原来src中的值放到新的3行5列中上述位置；而每个位置对应的值为原来src的值
$$1=(0,0);2=(1,1);3=(2,2);4=(0,3)$$
那么全零的矩阵变化如下
$$\left[\begin{array}{ccccc}0& 0& 0& 0& 0\\ 0& 0& 0& 0& 0\\ 0& 0& 0& 0& 0\end{array}\right]\to \left[\begin{array}{ccccc}1& 0& 0& 4& 0\\ 0& 2& 0& 0& 0\\ 0& 0& 3& 0& 0\end{array}\right]$$

• dim=1时分析；那么我们就可以得到dim=1为指定的Index，其他的保持不变，以前多少行就多少行
• 原始张量src
  $$src=torch.tensor([[1,2,3,4,5],[6,7,8,9,10]])$$
• 索引张量index:
  $$index=torch.tensor([[0,1,2],[0,1,4]])$$
  $$[[0,1,2],[0,1,4]]->[[(0,0),(0,1),(0,2)],[(1,0),(1,1),(1,4)]]$$
  也就是说将原来src的值放到新的3行5列矩阵上述位置，如果没有坐标的默认为0
  $$1=(0,0);2=(0,1);3=(0,2);4=放弃;5=放弃$$
  $$6=(1,0);7=(1,1);8=(1,4);9=放弃;10=放弃$$
  代入到新的矩阵如下：
  $$\left[\begin{array}{ccccc}0& 0& 0& 0& 0\\ 0& 0& 0& 0& 0\\ 0& 0& 0& 0& 0\end{array}\right]\to \left[\begin{array}{ccccc}1& 2& 3& 0& 0\\ 6& 7& 0& 0& 8\\ 0& 0& 0& 0& 0\end{array}\right]$$
• 代码如下：

import torch

src = torch.arange(1, 11).reshape(2, 5)
# src=tensor([[ 1,  2,  3,  4,  5],
#         [ 6,  7,  8,  9, 10]])
index = torch.tensor([[0, 1, 2], [0, 1, 4]])
# index=tensor([[0, 1, 2],
#         [0, 1, 4]])
output = torch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src)
# output=tensor([[1, 2, 3, 0, 0],
#         [6, 7, 0, 0, 8],
#         [0, 0, 0, 0, 0]])
print(f"src={src}")
print(f"index={index}")
print(f"output={output}")

以上手动计算的跟程序结果一致！

6. torch.squeeze

得到一个新的张量，它将输入的张量中的大小为1的维度去除掉
如果输入的张量大小为input_shape=(A×1×B×C×1×D);输出为output_shape=(A×B×C)

import torch

input = torch.ones(2,1,1,4)
input_squeeze = torch.squeeze(input)
print(f"input={input}")
print(f"input.shape={input.shape}")
# input.shape=torch.Size([2, 1, 1, 4])
print(f"input_squeeze={input_squeeze}")
print(f"input_squeeze.shape={input_squeeze.shape}")
# input_squeeze.shape=torch.Size([2, 4])

7. torch.stack

沿着一个新的维度连接张量序列，dim 可以指定沿着指定维度堆叠，注意，是新生成一个维度放到指定为dim上

a = torch.ones(2,3,4)
b = torch.zeros(2,3,4)
a_b_dim_0 = torch.stack((a,b),dim=0)
a_b_dim_1 = torch.stack((a,b),dim=1)
a_b_dim_2 = torch.stack((a,b),dim=2)
print(f"a_b_dim_0={a_b_dim_0}")
print(f"a_b_dim_0.shape={a_b_dim_0.shape}")
print(f"a_b_dim_1={a_b_dim_1}")
print(f"a_b_dim_1.shape={a_b_dim_1.shape}")
print(f"a_b_dim_2={a_b_dim_2}")
print(f"a_b_dim_2.shape={a_b_dim_2.shape}")

a_b_dim_0=tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]],

         [[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]]],


        [[[0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.]],

         [[0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.]]]])
a_b_dim_0.shape=torch.Size([2, 2, 3, 4])
a_b_dim_1=tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]],

         [[0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.]]],


        [[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]],

         [[0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.]]]])
a_b_dim_1.shape=torch.Size([2, 2, 3, 4])
a_b_dim_2=tensor([[[[1., 1., 1., 1.],
          [0., 0., 0., 0.]],

         [[1., 1., 1., 1.],
          [0., 0., 0., 0.]],

         [[1., 1., 1., 1.],
          [0., 0., 0., 0.]]],


        [[[1., 1., 1., 1.],
          [0., 0., 0., 0.]],

         [[1., 1., 1., 1.],
          [0., 0., 0., 0.]],

         [[1., 1., 1., 1.],
          [0., 0., 0., 0.]]]])
a_b_dim_2.shape=torch.Size([2, 3, 2, 4])