pytorch tensot unsqueeze code example
Example 1: pytorch tensor add one dimension
# ADD ONE DIMENSION: .unsqueeze(dim)
my_tensor = torch.tensor([1,3,4])
# tensor([1,3,4])
my_tensor.unsqueeze(0)
# tensor([[1,3,4]])
my_tensor.unsqueeze(1)
# tensor([[1],
# [3],
# [4]])
Example 2: pytorch squeeze
x = torch.zeros(2, 1, 2, 1, 2)
x.size()
>>> torch.Size([2, 1, 2, 1, 2])
y = torch.squeeze(x) # remove 1
y.size()
>>> torch.Size([2, 2, 2])
y = torch.squeeze(x, 0)
y.size()
>>> torch.Size([2, 1, 2, 1, 2])
y = torch.squeeze(x, 1)
y.size()
>>> torch.Size([2, 2, 1, 2])