numpy.squeeze() 函数使用

numpy.squeeze(a, axis=None)，axis的取值可为None 或 int 或 tuple of ints

在数组 shape 中删除一维项。

1
2
3
4
5
6
7
8
9

import numpy as np
import torch

x = torch.randn(1, 1, 3, 1, 2, 1)
print(x)
print(x.shape)
y = np.squeeze(x)
print(y)
print(y.shape)

在不指定 axis 的值时，函数会删除所有一维项，输出如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

tensor([[[[[[ 1.0555],
[-0.1646]]],

[[[ 1.8298],
[ 0.9088]]],

[[[ 0.8458],
[ 0.2796]]]]]])
torch.Size([1, 1, 3, 1, 2, 1])
tensor([[ 1.0555, -0.1646],
[ 1.8298, 0.9088],
[ 0.8458, 0.2796]])
torch.Size([3, 2])

删除指定维度（删除第四维项）：

1
2
3
4

x = torch.randn(1, 1, 3, 1, 2, 1)
y = np.squeeze(x, 3)
print(y)
print(y.shape)

输出：

1
2
3
4
5
6
7
8
9

tensor([[[[[-1.0309],
[ 1.1256]],

[[ 0.3259],
[ 2.2576]],

[[ 0.1139],
[-0.0633]]]]])
torch.Size([1, 1, 3, 2, 1])

将数据转化为 numpy 数据，可以使用 tuple 类型

1
2
3
4
5

x = torch.randn(1, 1, 3, 1, 2, 1)
print(x.shape)
x = np.array(x)
y = np.squeeze(x, (0, 1))
print(y.shape)

输出：

1 2	torch.Size([1, 1, 3, 1, 2, 1]) (3, 1, 2, 1)

当我们删除非一维项时，此时程序不会报错，也不会改变数据 shape ：

1
2
3
4

x = torch.randn(1, 1, 3, 1, 2, 1)
y = np.squeeze(x, 2)
print(y)
print(y.shape)

输出：

1
2
3
4
5
6
7
8
9
10
11

tensor([[[[[[ 0.3623],
[-0.0908]]],

[[[-1.6641],
[ 0.4365]]],

[[[ 0.3317],
[ 0.8052]]]]]])
torch.Size([1, 1, 3, 1, 2, 1])

官方文档 numpy.squeeze