torch.stack 理解

import numpy as np
import torch
arr1 = np.array([[1, 2, 3], [4, 5, 6]])
arr2 = np.array([[7, 8, 9], [5, 4, 3]])
c = torch.from_numpy(arr1)
d = torch.from_numpy(arr2)


t_stack2 = torch.stack([c,d,d], dim=2)
print(t_stack2)
# 在第2维上进行数据拼接
tensor([[[1, 7, 7],
         [2, 8, 8],
         [3, 9, 9]],

        [[4, 5, 5],
         [5, 4, 4],
         [6, 3, 3]]])


t_stack0 = torch.stack([c,d,d], dim=0)
print(t_stack0)
# 在第0维上进行数据拼接，原来的第0维变成第二维，原来的第1维变成第2维
tensor([[[1, 2, 3],
         [4, 5, 6]],

        [[7, 8, 9],
         [5, 4, 3]],

        [[7, 8, 9],
         [5, 4, 3]]])


t_stack1 = torch.stack([c,d,d], dim=1)
print(t_stack1)
# 在第1维上进行数据拼接，原来的第1维变成第2维
tensor([[[1, 2, 3],
         [7, 8, 9],
         [7, 8, 9]],

        [[4, 5, 6],
         [5, 4, 3],
         [5, 4, 3]]])