torch.stack()简要介绍

定义

torch.stack((a,b),dim)

用法

stack操作后会增加维度数量，在dim维度上对a和b进行堆叠

示例

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)
 
>>tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
>>tensor([[11, 22, 33],
        [44, 55, 66],
        [77, 88, 99]])

在dim=0上堆叠a和b

torch.stack((a,b),dim=0)
 
>>tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],
 
        [[11, 22, 33],
         [44, 55, 66],
         [77, 88, 99]]])

在dim=1上堆叠a和b

torch.stack((a,b),dim=1)
 
>>tensor([[[ 1,  2,  3],
         [11, 22, 33]],
 
        [[ 4,  5,  6],
         [44, 55, 66]],
 
        [[ 7,  8,  9],
         [77, 88, 99]]])

在dim=2上堆叠a和b

torch.stack((a,b),dim=2)
 
>>tensor([[[ 1, 11],
         [ 2, 22],
         [ 3, 33]],
 
        [[ 4, 44],
         [ 5, 55],
         [ 6, 66]],
 
        [[ 7, 77],
         [ 8, 88],
         [ 9, 99]]])