【深度學習】PyTorch之Squeeze()和Unsqueeze()

1. unsqueeze()

該函數用來增加某個維度。在PyTorch中維度是從0開始的。

import torch

a = torch.arange(0, 9)
print(a)

結果：

tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])

利用view()改變tensor的形狀。值得注意的是view不會修改自身的數據，返回的新tensor與源tensor共享內存；同時必須保證前后元素總數一致。

a = a.view(3, 3)
print(f"a:{a} \n shape:{a.shape}")

結果：

a:tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]]) 
 shape:torch.Size([3, 3])

在第一個維度(即維度序號為0)前增加一個維度。

a = a.unsqueeze(0)
print(f"a:{a}\nshape:{a.shape}")

結果：

a:tensor([[[0, 1, 2],
         [3, 4, 5],
         [6, 7, 8]]])
shape:torch.Size([1, 3, 3])

同理，可在其他位置添加維度，在這里就不舉例了。

2. squeeze()

該函數用來減少某個維度。

print(f"1.   a:{a}\nshape:{a.shape}")
a = a.unsqueeze(0)
a = a.unsqueeze(2)
print(f"2.   a:{a}\nshape:{a.shape}")
a = a.squeeze(2)
print(f"3.   a:{a}\nshape:{a.shape}")

結果：

1.   a:tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
shape:torch.Size([3, 3])
2.   a:tensor([[[[0, 1, 2]],

         [[3, 4, 5]],

         [[6, 7, 8]]]])
shape:torch.Size([1, 3, 1, 3])
3.   a:tensor([[[0, 1, 2],
         [3, 4, 5],
         [6, 7, 8]]])
shape:torch.Size([1, 3, 3])

3. 下面是運用上述兩個函數，並進行一次卷積的例子。

from torchvision.transforms import  ToTensor
import torch as t
from torch import nnimport cv2
import numpy as np
import cv2

to_tensor = ToTensor()
# 加載圖像
lena = cv2.imread('lena.jpg', cv2.IMREAD_GRAYSCALE)
cv2.imshow('lena', lena)
# input = to_tensor(lena) 將ndarray轉換為tensor，自動將[0,255]歸一化至[0,1]。
input = to_tensor(lena).unsqueeze(0)
# 初始化卷積參數
kernel = t.ones(1, 1, 3, 3)/-9
kernel[:, :, 1, 1] = 1
conv = nn.Conv2d(1, 1, 3, 1, padding=1, bias=False)
conv.weight.data = kernel.view(1, 1, 3, 3)
# 輸出
out = conv(input)
out = out.squeeze(0)
print(out.shape)
out = out.unsqueeze(3)
print(out.shape)
out = out.squeeze(0)
print(out.shape)
out = out.detach().numpy()
# 縮放到0~最大值
cv2.normalize(out, out, 1.0, 0, cv2.NORM_INF)
cv2.imshow("lena-result", out)
cv2.waitKey()

結果：

         

torch.Size([1, 304, 304])
torch.Size([1, 304, 304, 1])
torch.Size([304, 304, 1])
<class 'numpy.ndarray'> (304, 304, 1)

 

參考文獻

[1] 陳雲.深度學習框架之PyTorch入門與實踐.北京：電子工業出版社，2018.