1. 程式人生 > 其它 >torch.squeeze() 和torch.unsqueeze()的用法,給tensor刪除或者新增維度為1的維度(squeeze和unsqueeze)

torch.squeeze() 和torch.unsqueeze()的用法,給tensor刪除或者新增維度為1的維度(squeeze和unsqueeze)

import torch
x = torch.randn(3,1,4,1,2)
a = x.squeeze(dim=1)  # 成員函式刪除第二維度
print(a.shape)
b = torch.squeeze(x, dim=1)  # torch自帶函式刪除第二維度
print(b.shape)
c = torch.squeeze(x, dim=3)  # 刪除第三維度
print(c.shape)
d = torch.squeeze(x)  # 若不標註刪除第幾維度,則會刪除所有為1的維度
print(d.shape)

輸出:
torch.Size([3, 4, 1, 2])
torch.Size([3, 4, 1, 2])
torch.Size([3, 1, 4, 2])
torch.Size([3, 4, 2])

注意:squeeze()函式只能刪除維度為1的維度,若不為1,該操作無效,但不會報錯

b = torch.unsqueeze(a, dim=1)
print(b.shape)

輸出為:
torch.Size([3, 1, 4, 1, 2])

 

在pytorch中, 用自帶的torch.unsqueeze()和tensor的成員函式unsqueeze()可以為tensor新增維度為1的維度

e = torch.unsqueeze(x, dim=0) # 在第一維度新增維度
print(e.shape)
f = x.unsqueeze(dim=0)
print(f.shape)

輸出:
torch.Size([1, 3, 1, 4, 1, 2])
torch.Size([1, 3, 1, 4, 1, 2])

 

先看torch.squeeze() 這個函式主要對資料的維度進行壓縮,去掉維數為1的的維度,比如是一行或者一列這種,一個一行三列(1,3)的數去掉第一個維數為一的維度之後就變成(3)行。squeeze(a)就是將a中所有為1的維度刪掉。不為1的維度沒有影響。a.squeeze(N) 就是去掉a中指定的維數為一的維度。還有一種形式就是b=torch.squeeze(a,N) a中去掉指定的定的維數為一的維度。

再看torch.unsqueeze()這個函式主要是對資料維度進行擴充。給指定位置加上維數為一的維度,比如原本有個三行的資料(3),在0的位置加了一維就變成一行三列(1,3)。a.squeeze(N) 就是在a中指定位置N加上一個維數為1的維度。還有一種形式就是b=torch.squeeze(a,N) a就是在a中指定位置N加上一個維數為1的維度