torch中squeeze与unsqueeze用法

时间:2020-11-30 21:32:22

import torch

torch中的squeeze与unsqueeze作用是去除/添加维度为1的行

例如,a=torch.randn(2,3)

那么b=a.unsqueeze(0),b为(1,2,3)的矩阵

类似的,b=a.unsqueeze(1),b为(2,1,3)的矩阵

b=a.unsqueeze(2),b为(2,3,1)的矩阵

同理,如果a=torch.randn(2,1,3)

则b=a.squeeze(2)=torch.squeeze(a),b为(2,3)的矩阵。