torch.unsqueeze、np.expand_dims详解

超链接：深度学习工作常用方法汇总，矩阵维度变化、图片、视频等操作，包含（torch、numpy、opencv等）

---

增加维度 ： unsqueeze、np.expand_dims

torch版：

x.unsqueeze(dim=0)

简介：

将矩阵x在dim=0维度上增加一个维度。 [3, 3, 2] -->>-- [1, 3, 3, 2]
可以理解为，在dim=n的前面增加一个维度。索引从0开始

用途之一：torch 模型训练中dataloader函数返回的数据为[batch, 3, 224, 224]大小的数据，会增加一个batch维度，但是我们在预测的时候，一般都是一张图片直接进入模型，进行预测，可是模型的输入需要batch维度，我们为了维度对应，一般都用 增加维度的方法实现。

torch示例：

import torch
import numpy as np

x = torch.rand((2, 2, 3, 3))
b = x.unsqueeze(0)
c = x.unsqueeze(1)
d = x.unsqueeze(2)
print('x_shape:', x.shape)  # torch.Size([2, 2, 3, 3])
print('b_shape:', b.shape)  # b_shape: torch.Size([1, 2, 2, 3, 3])
print('c_shape:', c.shape)  # c_shape: torch.Size([2, 1, 2, 3, 3])
print('d_shape:', d.shape)  # d_shape: torch.Size([2, 2, 1, 3, 3])

numpy版：

np.expand_dims(arr, 0)

numpy示例：

import torch
import numpy as np

x = np.array(([1, 2], [3, 4]))
b = np.expand_dims(x, axis=0)
c = np.expand_dims(x, axis=1)
print('x_shape:', x.shape)  # (2, 2)
print('b_shape:', b.shape)  # b_shape: (1, 2, 2)
print('c_shape:', c.shape)  # c_shape: (2, 1, 2)