【深度学习】PyTorch之Squeeze和Unsqueeze

1. unsqueeze)

该函数用来增加某个维度。在PyTorch中维度是从0开始的。

import torch

a = torch.arange0, 9)
printa)

结果：

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

利用view)改变tensor的形状。值得注意的是view不会修改自身的数据，返回的新tensor与源tensor共享内存；同时必须保证前后元素总数一致。

a = a.view3, 3)
printf"a:{a} 
 shape:{a.shape}")

结果：

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

在第一个维度即维度序号为0)前增加一个维度。

a = a.unsqueeze0)
printf"a:{a}
shape:{a.shape}")

结果：

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

同理，可在其他位置添加维度，在这里就不举例了。

2. squeeze)

该函数用来减少某个维度。

printf"1.   a:{a}
shape:{a.shape}")
a = a.unsqueeze0)
a = a.unsqueeze2)
printf"2.   a:{a}
shape:{a.shape}")
a = a.squeeze2)
printf"3.   a:{a}
shape:{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_tensorlena) 将ndarray转换为tensor，自动将[0,255]归一化至[0,1]。
input = to_tensorlena).unsqueeze0)
# 初始化卷积参数
kernel = t.ones1, 1, 3, 3)/-9
kernel[:, :, 1, 1] = 1
conv = nn.Conv2d1, 1, 3, 1, padding=1, bias=False)
conv.weight.data = kernel.view1, 1, 3, 3)
# 输出
out = convinput)
out = out.squeeze0)
printout.shape)
out = out.unsqueeze3)
printout.shape)
out = out.squeeze0)
printout.shape)
out = out.detach).numpy)
# 缩放到0~最大值
cv2.normalizeout, 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.