pytorch中torch.stack用法虽简单,但不好理解-参数列表

时间:2024-04-01 16:10:10

tensors :为一系列输入张量,类型为turple和List
dim :新增维度的(下标)位置,当dim = -1时默认最后一个维度;范围必须介于 0 到输入张量的维数之间,默认是dim=0,在第0维进行连接
返回值:输出新增维度后的张量

情况一:输入数据为1维数据

dim = 0 : 在第0维进行连接,相当于在行上进行组合(输入张量为一维,输出张量为两维)

import torch

a = torch.tensor([1, 2, 3])

b = torch.tensor([11, 22, 33])

#在第0维进行连接,相当于在行上进行组合,取a的一行,b的一行,构成一个新的tensor(输入张量为一维,输出张量为两维)

c = torch.stack([a, b],dim=0)          

print(a)

print(b)

print(c.size())

print(c)

输出:
tensor([1, 2, 3])
tensor([11, 22, 33])
torch.Size([2, 3])
tensor([[ 1,  2,  3],
        [11, 22, 33]])

dim = 1 :在第1维进行连接,相当于在对应行上面对列元素进行组合(输入张量为一维,输出张量为两维)

import torch

a = torch.tensor([1, 2, 3])

b = torch.tensor([11, 22, 33])

print(a)

print(b)

#在第1维进行连接,相当于在对应行上面对列元素进行组合,取a的一列,b的一列,构成新的tensor的一行(输入张量为一维,输出张量为两维)

c = torch.stack([a, b],dim=1)

print(c.size())

print(c)

输出:
tensor([1, 2, 3])
tensor([11, 22, 33])
torch.Size([3, 2])
tensor([[ 1, 11],
        [ 2, 22],
        [ 3, 33]])

情况二:输入数据为2维数据

dim=0:表示在第0维进行连接,相当于在通道维度上进行组合(输入张量为两维,输出张量为三维),注意:此处输入张量维度为二维,因此dim最大只能为2。

import torch

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

b = torch.tensor([[11, 22, 33], [44, 55, 66], [77, 88, 99]])

print(a)

print(b)
#在第0维进行连接,相当于在通道维度上进行组合
#即取a的所有数据,作为新tensor的一个分量
#取b的所有数据,作为新tensor的另一个分量
#(输入张量为两维,输出张量为三维)

c = torch.stack([a, b],dim=0)

print(c.size())

print(c)

输出:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
tensor([[11, 22, 33],
        [44, 55, 66],
        [77, 88, 99]])
torch.Size([2, 3, 3])
tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[11, 22, 33],
         [44, 55, 66],
         [77, 88, 99]]])

dim=1:表示在第1维进行连接,相当于对相应通道中每个行进行组合,注意:此处输入张量维度为二维,因此dim最大只能为2。

import torch

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

b = torch.tensor([[11, 22, 33], [44, 55, 66], [77, 88, 99]])

print(a)

print(b)

#在第1维(行)进行连接,相当于对相应通道中每个行进行组合
#取a的一行,b的一行,作为新tensor的第1行和第2行
#原来a:3*3,b:3*3,新tensor:3*2*3

c = torch.stack([a, b], 1)

print(c.size())

print(c)

输出:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
tensor([[11, 22, 33],
        [44, 55, 66],
        [77, 88, 99]])
torch.Size([3, 2, 3])
tensor([[[ 1,  2,  3],
         [11, 22, 33]],

        [[ 4,  5,  6],
         [44, 55, 66]],

        [[ 7,  8,  9],
         [77, 88, 99]]])

dim=2:表示在第2维进行连接,相当于对相应行中每个列元素进行组合,注意:此处输入张量维度为二维,因此dim最大只能为2。

import torch

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

b = torch.tensor([[11, 22, 33], [44, 55, 66], [77, 88, 99]])

print(a)

print(b)

#在第2维进行连接,相当于对相应行中每个列元素进行组合
#针对每行,取a、b的第一列数据,构成tensor的第一行
#针对每行,取a、b的第二列数据,构成tensor的第二行
#,针对每行取a、b的第三列数据,构成tensor的第三行
#原来a:3*3,b:3*3,新tensor:3*3*2
c = torch.stack([a, b], 2)

print(c.size())

print(c)

输出:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
tensor([[11, 22, 33],
        [44, 55, 66],
        [77, 88, 99]])
torch.Size([3, 3, 2])
tensor([[[ 1, 11],
         [ 2, 22],
         [ 3, 33]],

        [[ 4, 44],
         [ 5, 55],
         [ 6, 66]],

        [[ 7, 77],
         [ 8, 88],
         [ 9, 99]]])

情况三:输入数据为3维数据

dim=0:表示在第0维进行连接,相当于在通道维进行拼接。注意:此处输入张量维度为三维,因此dim最大只能为3。

import torch

a = torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],[[10, 20, 30], [40, 50, 60], [70, 80, 90]]])

b = torch.tensor([[[11, 22, 33], [44, 55, 66], [77, 88, 99]], [[110, 220, 330], [440, 550, 660], [770, 880, 990]]])

print(a)

print(b)

#表示在第0维进行连接,取整个a作为新tensor的一个分量,取整个b作为新tensor的一个分量
c = torch.stack([a, b], 0)

print(c)

输出:
tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 20, 30],
         [40, 50, 60],
         [70, 80, 90]]])
torch.Size([2, 3, 3])
tensor([[[ 11,  22,  33],
         [ 44,  55,  66],
         [ 77,  88,  99]],

        [[110, 220, 330],
         [440, 550, 660],
         [770, 880, 990]]])
torch.Size([2, 3, 3])
torch.Size([2, 2, 3, 3])
tensor([[[[  1,   2,   3],
          [  4,   5,   6],
          [  7,   8,   9]],

         [[ 10,  20,  30],
          [ 40,  50,  60],
          [ 70,  80,  90]]],


        [[[ 11,  22,  33],
          [ 44,  55,  66],
          [ 77,  88,  99]],

         [[110, 220, 330],
          [440, 550, 660],
          [770, 880, 990]]]])

dim=1:表示在第1维进行连接,取各自的第1维度数据,进行拼接。注意:此处输入张量维度为三维,因此dim最大只能为3。 

import torch

a = torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],[[10, 20, 30], [40, 50, 60], [70, 80, 90]]])

b = torch.tensor([[[11, 22, 33], [44, 55, 66], [77, 88, 99]], [[110, 220, 330], [440, 550, 660], [770, 880, 990]]])

print(a)
print(a.size())

print(b)
print(b.size())

#表示在第1维进行连接,取a的第一维数据[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
#取b的第一维数据[[11, 22, 33], [44, 55, 66], [77, 88, 99]]作为新tensor的一个分量

#取a的第一维数据[[10, 20, 30], [40, 50, 60], [70, 80, 90]]
#取b的第一维数据[[110, 220, 330], [440, 550, 660], [770, 880, 990]]作为新tensor的另一个分量
c = torch.stack([a, b], 1)

print(c.size())

print(c)

输出:
tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 20, 30],
         [40, 50, 60],
         [70, 80, 90]]])
torch.Size([2, 3, 3])
tensor([[[ 11,  22,  33],
         [ 44,  55,  66],
         [ 77,  88,  99]],

        [[110, 220, 330],
         [440, 550, 660],
         [770, 880, 990]]])
torch.Size([2, 3, 3])
torch.Size([2, 2, 3, 3])
tensor([[[[  1,   2,   3],
          [  4,   5,   6],
          [  7,   8,   9]],

         [[ 11,  22,  33],
          [ 44,  55,  66],
          [ 77,  88,  99]]],


        [[[ 10,  20,  30],
          [ 40,  50,  60],
          [ 70,  80,  90]],

         [[110, 220, 330],
          [440, 550, 660],
          [770, 880, 990]]]])

dim=2:表示在第2维进行连接,取各自的第2维度数据,进行拼接。注意:此处输入张量维度为三维,因此dim最大只能为3。

import torch

a = torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],[[10, 20, 30], [40, 50, 60], [70, 80, 90]]])

b = torch.tensor([[[11, 22, 33], [44, 55, 66], [77, 88, 99]], [[110, 220, 330], [440, 550, 660], [770, 880, 990]]])

print(a)
print(a.size())

print(b)
print(b.size())

#表示在第1维进行连接,取a的第2维数据[1, 2, 3]
#取b的第2维数据[11, 22, 33]作为新tensor的一个分量

#取a的第2维数据[4, 5, 6]
#取b的第2维数据[44, 55, 66]作为新tensor的一个分量

#取a的第2维数据[4, 5, 6]
#取b的第2维数据[44, 55, 66]作为新tensor的一个分量

#取a的第2维数据[7, 8, 9]
#取b的第2维数据[77, 88, 99]作为新tensor的一个分量

#取a的第2维数据[10, 20, 30]
#取b的第2维数据[110, 220, 330]作为新tensor的一个分量

#取a的第2维数据[40, 50, 60]
#取b的第2维数据[440, 550, 660]作为新tensor的一个分量

#取a的第2维数据[70, 80, 90]
#取b的第2维数据[770, 880, 990]作为新tensor的一个分量
c = torch.stack([a, b], 2)
print(c.size())

print(c)

输出:
tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 20, 30],
         [40, 50, 60],
         [70, 80, 90]]])
torch.Size([2, 3, 3])
tensor([[[ 11,  22,  33],
         [ 44,  55,  66],
         [ 77,  88,  99]],

        [[110, 220, 330],
         [440, 550, 660],
         [770, 880, 990]]])
torch.Size([2, 3, 3])
torch.Size([2, 3, 2, 3])
tensor([[[[  1,   2,   3],
          [ 11,  22,  33]],

         [[  4,   5,   6],
          [ 44,  55,  66]],

         [[  7,   8,   9],
          [ 77,  88,  99]]],


        [[[ 10,  20,  30],
          [110, 220, 330]],

         [[ 40,  50,  60],
          [440, 550, 660]],

         [[ 70,  80,  90],
          [770, 880, 990]]]])

dim=3:表示在第3维进行连接,取各自的第3维度数据,进行拼接。注意:此处输入张量维度为三维,因此dim最大只能为3。

import torch

a = torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],[[10, 20, 30], [40, 50, 60], [70, 80, 90]]])

b = torch.tensor([[[11, 22, 33], [44, 55, 66], [77, 88, 99]], [[110, 220, 330], [440, 550, 660], [770, 880, 990]]])

print(a)
print(a.size())

print(b)
print(b.size())

#针对第二维数据,在每个第二维度相同的情况下,取各自的列数据,构成新tensor的一行
c = torch.stack([a, b], 3)
print(c.size())

print(c)

输出:
tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 20, 30],
         [40, 50, 60],
         [70, 80, 90]]])
torch.Size([2, 3, 3])
tensor([[[ 11,  22,  33],
         [ 44,  55,  66],
         [ 77,  88,  99]],

        [[110, 220, 330],
         [440, 550, 660],
         [770, 880, 990]]])
torch.Size([2, 3, 3])
torch.Size([2, 3, 3, 2])
tensor([[[[  1,  11],
          [  2,  22],
          [  3,  33]],

         [[  4,  44],
          [  5,  55],
          [  6,  66]],

         [[  7,  77],
          [  8,  88],
          [  9,  99]]],


        [[[ 10, 110],
          [ 20, 220],
          [ 30, 330]],

         [[ 40, 440],
          [ 50, 550],
          [ 60, 660]],

         [[ 70, 770],
          [ 80, 880],
          [ 90, 990]]]])

总结:m个序列数据,在某个维度k进行拼接,该维度大小为n,则拼接后形成了*n*m*大小,具体拼接过程是取m个序列数据,k-1维(设k-1维大小为x,从x=1开始取)相同情况下的第1个数据,构成新tensor的一个行;第二个数据...,第三个数据...构成tensor的新行;然后从x=2开始执行同样的操作