[pytorch] list类型的batch转tensor,cat与stack的区别

有一个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]]])


版权声明:本文为tagagi原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接和本声明。