有一个list,其中的元素为一个tensor,list的长度即为batch_size:
a = torch.randn(2, 3)
b = torch.randn(2, 3)
c = torch.randn(2, 3)
l = [a, b ,c]
需要将其转换为一个大的tensor作为模型的输入,二种方法
1) 使用stack,注意stack和后续的cat一样,返回一个心得tensor,不修改原来的
torch.stack(l, 0)
Out[8]:
tensor([[[ 0.2798, -0.6800, -0.0358],
[ 1.4812, -1.0230, -0.3167]],
[[ 2.9450, -1.5106, -0.2924],
[ 0.6985, -1.7721, -0.0210]],
[[-1.1184, 1.2074, -1.2872],
[ 1.2916, 1.3292, 1.3636]]])
2)cat + unsqeeze。需要注意的是,stack是堆叠,因而会直接在原来每个tensor上扩展一个维度,而cat是拼接,需要自己使用unsqueeze进行手动扩展维度。
torch.cat([i.unsqueeze(dim=0) for i in l], dim=0)
Out[22]:
tensor([[[ 0.2798, -0.6800, -0.0358],
[ 1.4812, -1.0230, -0.3167]],
[[ 2.9450, -1.5106, -0.2924],
[ 0.6985, -1.7721, -0.0210]],
[[-1.1184, 1.2074, -1.2872],
[ 1.2916, 1.3292, 1.3636]]])