numpy 与 torch中压缩、扩展维度的方法

写了一段时间的torch，发现torch与numpy的函数非常像，使用起来非常方便。今天发现 numpy 与 torch中压缩、扩展维度的方法稍有不同。

torch 压缩、扩展维度：

1
2
3
4
5
6
7
8
9
10
11

# 扩展维度使用unsqueeze()
A = torch.ones(8,8)

A = A.unsqueeze(2)
A = torch.unsqueeze(A,2)

# 压缩维度使用squeeze()
A = torch.ones(8,8,1)

A = A.squeeze(-1)
A = torch.squeeze(A,-1)

可以发现torch扩展维度，压缩维度非常方便，今天使用numpy想做形状扩展时，发现使用unsqueeze()，直接报错。

numpy 压缩、扩展维度：

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

import numpy as np

# numpy扩展维度，使用expand_dims()
A = np.ones([8,8])

A = A.expand_dims(2) #直接报错，numpy.array没有expand_dims函数
A = np.expand_dims(A,2)

# numpy压缩维度，使用squeeze()
A = np.ones([8,8,1])

A = A.squeeze(-1)
A = np.squeeze(A,-1) #两个方法都能使用