torch.squeeze:将tensor中大小为1的维度删除
torch.squeeze(tensor,dim)

代码示例:

import torch
a=torch.ones(2,2,2,1,1)
b=torch.squeeze(a)
c=torch.squeeze(a,0)
d=torch.squeeze(a,3)
print("a的shape:",a.shape)
print("b的shape:",b.shape)
print("c的shape:",c.shape)
print("d的shape:",d.shape)

输出:

a的shape: torch.Size([2, 2, 2, 1, 1])
b的shape: torch.Size([2, 2, 2])
c的shape: torch.Size([2, 2, 2, 1, 1])
d的shape: torch.Size([2, 2, 2, 1])

注:对特定维度进行squeeze,如果该维度大小不为1,则保持原来的shape不变,如代码中的c的shape。

Logo

华为开发者空间,是为全球开发者打造的专属开发空间,汇聚了华为优质开发资源及工具,致力于让每一位开发者拥有一台云主机,基于华为根生态开发、创新。

更多推荐