torch.stack([tensor1, tensor2, tensor3…], dim=0)
基本作用: 将几个矩阵按照dim指定维度堆叠在一起;具体如何堆叠看下方例子
import torch
x = torch.rand((2,3))
y = torch.rand((2,3))
z = torch.stack([x, y], dim=0)
z2 = torch.stack([x, y], dim=1)
z3 = torch.stack([x, y], dim=2)
print("x:",x)
print("y:", y)
print("z:", z)
print("z2:",z2)
print("z3:", z3)
输出:
x: tensor([[0.2380, 0.3537, 0.8676],
[0.6321, 0.8611, 0.1720]])
y: tensor([[0.0692, 0.4968, 0.2710],
[0.2103, 0.6017, 0.3630]])
z: tensor([[[0.2380, 0.3537, 0.8676],
[0.6321, 0.8611, 0.1720]],
[[0.0692, 0.4968, 0.2710],
[0.2103, 0.6017, 0.3630]]])
z2: tensor([[[0.2380, 0.3537, 0.8676],
[0.0692, 0.4968, 0.2710]],
[[0.6321, 0.8611, 0.1720],
[0.2103, 0.6017, 0.3630]]])
z3: tensor([[[0.2380, 0.0692],
[0.3537, 0.4968],
[0.8676, 0.2710]],
[[0.6321, 0.2103],
[0.8611, 0.6017],
[0.1720, 0.3630]]])
[Finished in 1.5s]
总结
假设list中tensor的个数为n, 每个tensor的shape为[w, h]
- dim = 0 就是将list中几个矩阵直接按顺序堆叠起来,堆叠结果的shape为[n, w, h]
- dim = 1 是将list中每个矩阵的的每行拿出来重新组成w个矩阵,然后堆叠在一起,最后结果shape为[w, n, h]; 具体来说就是将每个tensor的第一行拿出来组成一个tensor1(shape=[n, h]), 然后每个tensor第二行拿出来组成tensor2, …依次进行下去,最后将tensor1, tensor2…堆叠起来
- dim = 2 将每个tensor的每行拿出来组成tenor, 然后转置,得到tensor_i, 最后将tensor_1, tensor2, …tensor_i堆叠起来, 最后结果的shape=[w, h, n]。具体来说,就先将每个tensor第一行拿出来组成矩阵再转置得到tensor1, 然后取每个tensor的第二行做相同操作,得到tensor2,…,最后将tensor1, tensor2,…全部堆叠起来