Tensor的组合与分块-02

Tensor的组合与分块

　　组合操作是指将不同的Tensor叠加起来， 主要有torch.cat()和torch.stack()两个函数。 cat即concatenate的意思， 是指沿着已有的数据的某一维度进行拼接， 操作后数据的总维数不变， 在进行拼接时， 除了拼接的维度之外， 其他维度必须相同。 而torch.stack()函数指新增维度， 并按照指定的维度进行叠加，

import torch 

# 创建两个2×2的Tensor
a = torch.Tensor([[1,2],[3,4]])
print(a,a.shape)

b = torch.Tensor([[5,6],[7,8]])
print(b,b.shape)

# 以第一维进行拼接， 则变成4×2的矩阵
c = torch.cat([a,b],0)
print(c,c.shape)

# 以第二维进行拼接， 则变成2*4的矩阵
d = torch.cat([a,b],1)
print(d,d.size())

结果输出：

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

import torch 

# 创建两个2×2的Tensor
a = torch.Tensor([[1,2],[3,4]])
print(a,a.shape)

>>   tensor([[1., 2.],
            [3., 4.]]) torch.Size([2, 2])

b = torch.Tensor([[5,6],[7,8]])
print(b,b.shape)

>>   tensor([[5., 6.],
            [7., 8.]]) torch.Size([2, 2])

# 以第0维进行stack， 叠加的基本单位为序列本身， 即a与b， 因此输出[a, b]， 输出维度为2×2×2
d=torch.stack([a,b],0)
print(d, d.size())
>>  tensor([[[1., 2.],
         [3., 4.]],

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

# 以第1维进行stack， 叠加的基本单位为每一行， 输出维度为2×2×2
e=torch.stack([a,b],1)
print(e, e.shape)

>> tensor([[[1., 2.],
         [5., 6.]],

        [[3., 4.],
         [7., 8.]]]) torch.Size([2, 2, 2])

# 以第2维进行stack， 叠加的基本单位为每一行的每一个元素， 输出维度为2×2×2
f=torch.stack([a,b],2)
print(f, f.shape)

>> tensor([[[1., 5.],
         [2., 6.]],

        [[3., 7.],
         [4., 8.]]]) torch.Size([2, 2, 2])

 　　分块则是与组合相反的操作， 指将Tensor分割成不同的子Tensor，主要有torch.chunk()与torch.split()两个函数， 前者需要指定分块的数量，而后者则需要指定每一块的大小， 以整型或者list来表示。 具体示例如下 ：

import torch 

a = torch.Tensor([[1,2,3], [4,5,6]])
print(a, a.size())
>> tensor([[1., 2., 3.],
        [4., 5., 6.]]) torch.Size([2, 3])

# 使用chunk， 沿着第0维进行分块， 一共分两块， 因此分割成两个1×3的Tensor
b = torch.chunk(a, 2, 0)
print(b)
>> (tensor([[1., 2., 3.]]), tensor([[4., 5., 6.]]))

# 沿着第1维进行分块， 因此分割成两个Tensor， 当不能整除时， 最后一个的维数会小于前面的
# 因此第一个Tensor为2×2， 第二个为2×1
c = torch.chunk(a, 2, 1)
print(c)
>> (tensor([[1., 2.],
        [4., 5.]]), tensor([[3.],
        [6.]]))

# 使用split， 沿着第0维分块， 每一块维度为2， 由于第一维维度总共为2， 因此相当于没有分割
d = torch.split(a, 2, 0)
print(d)
>> (tensor([[1., 2., 3.],
        [4., 5., 6.]]),)

# 沿着第1维分块， 每一块维度为2， 因此第一个Tensor为2×2， 第二个为2×1
e = torch.split(a, 2, 1)
print(e)
>> (tensor([[1., 2.],
        [4., 5.]]), tensor([[3.],
        [6.]]))
 
# split也可以根据输入的list进行自动分块， list中的元素代表了每一个块占的维度
f = torch.split(a, [1,2], 1)
print(f)
>> (tensor([[1.],
        [4.]]), tensor([[2., 3.],
        [5., 6.]]))