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

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

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

欢迎分享,转载请注明来源:内存溢出

原文地址: http://outofmemory.cn/langs/715550.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-04-25
下一篇 2022-04-25

发表评论

登录后才能评论

评论列表(0条)

保存