使用NumPy.squeeze轻松压缩多余的维度

一、NumPy.squeeze的作用

在TensorFlow、PyTorch、NumPy等深度学习库中，经常需要对张量（Tensor）数据进行处理。有时候，通过各种操作（比如切片、转置等）之后，可能会出现维度为1的情况，这时候可以使用NumPy.squeeze函数将张量中的长度为1的维度去掉。

import numpy as np

a = np.array([[[1, 2, 3]]])
print(a.shape)  # (1, 1, 3)

b = np.squeeze(a)
print(b.shape)  # (3,)

上面的代码中，使用NumPy的array函数创建了一个形状为(1, 1, 3)的三维张量。通过NumPy.squeeze函数，将长度为1的维度去掉，得到了形状为(3,)的一维数组。

二、使用NumPy.squeeze的优点

使用NumPy.squeeze函数的优点在于，它可以简化代码，提高代码的可读性。比如，在处理图像数据的时候，常常需要对张量进行转换，这时候可以使用squeeze函数去掉维度为1的轴。

import numpy as np

# 创建一个形状为(1, 28, 28, 1)的四维张量
x = np.random.randn(1, 28, 28, 1)

# 对张量进行切片操作，得到形状为(28, 28)的二维张量
y = x[0, :, :, 0]

# 使用NumPy.squeeze函数，去掉维度为1的轴
z = np.squeeze(y)

print(z.shape)  # (28, 28)

上面的代码中，使用np.random.randn函数创建了一个形状为(1, 28, 28, 1)的四维张量。然后，通过对张量x进行切片操作，得到了形状为(28, 28)的二维张量y。使用np.squeeze函数，去掉维度为1的轴，得到形状为(28, 28)的二维张量z。

三、NumPy.squeeze的注意事项

NumPy.squeeze函数虽然非常方便，但是在使用的时候需要注意一些细节问题。

首先，需要注意的是，函数会返回一个新数组，因此需要将结果赋值给一个新的变量。如果不这样做，原始数组不会被改变。

import numpy as np

# 创建一个形状为(1, 28, 28, 1)的四维张量
x = np.random.randn(1, 28, 28, 1)

# 使用NumPy.squeeze函数去掉维度为1的轴
np.squeeze(x)

print(x.shape)  # (1, 28, 28, 1)

上面的代码中，使用NumPy.squeeze函数去掉维度为1的轴，但是没有保存结果。因此，原始的数组x没有被改变。

其次，需要注意的是，如果要去掉的维度不是长度为1的维度，那么squeeze函数不会做任何改变。

import numpy as np

# 创建一个形状为(1, 3, 28, 28)的四维张量
x = np.random.randn(1, 3, 28, 28)

# 使用NumPy.squeeze函数去掉维度为1的轴
np.squeeze(x)

print(x.shape)  # (1, 3, 28, 28)

上面的代码中，创建了一个形状为(1, 3, 28, 28)的四维张量。虽然它的第一维长度为1，但是它不是长度为1的轴，因此squeeze函数不会做任何改变。

四、总结

NumPy.squeeze函数是一个非常方便的工具，可以用来轻松压缩多余的维度。在处理张量数据时，使用squeeze函数可以简化代码，提高代码的可读性。使用时需要注意，如果要去掉的维度不是长度为1的维度，那么squeeze函数不会做任何改变。