from einops import rearrange中rearrange操作

1. 安装einops包

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple einops
1

表示安装成功！!

2. rearrange用法

import torch
from einops import rearrange

a = torch.randn(1,2,3,2)  # 产生随机tensor，shape: torch.Size([1,2,3,2])
1
2
3
4

• rearrange按给出的模式(注释)重组张量，其中模式中字母只是个表示，没有具体含义

out1 = rearrange(a, 'b c h w -> b (c h w)')  # torch.Size([1,12])
out2 = rearrange(a, 'time c h w -> time (c h w)')  # torch.Size([1,12])
1
2

两个输出一摸一样！!

• 注释给定以后就代表当前维度，不能更改。如下

a = torch.randn(1,2,3,2)  # torch.Size([1,2,3,2]) 

out3 = rearrange(a, 'b c h w -> b (c h w)', c=2,h=3,w=2) # torch.Size([1,12])
# h,w互换
err1 = rearrange(a, 'b c h w -> b (c h w)', c=2,h=2,w=3) # 运行错误
# 强制将c=1
err2 = rearrange(a, 'b c h w -> b (c h w)', c=1,h=3,w=2) # 运行错误
1
2
3
4
5
6
7

err1输出：

err2输出：

• 将张量深度转换为广度(depth-to-space)

out1 = rearrange(a, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=1, w2=2) 
out2 = rearrange(a, 'b c (h h2) (w w2) -> b (h2 w2 c) h w', h2=1, w2=2)
'''
输出：
a:                                 
tensor([[[[ 0.7676, -0.4329],  
          [-1.2082,  1.1884],
          [-0.8015,  0.2629]],

         [[-1.3980, -0.0211],
          [-0.5598, -0.5385],
          [-0.6689, -0.0628]]]])
# torch.Size([1, 2, 3, 2])

out1:                              out2: 
tensor([[[[ 0.7676],                   tensor([[[[ 0.7676],
          [-1.2082],                           [-1.2082],
          [-0.8015]],                          [-0.8015]],

         [[-0.4329],                          [[-1.3980],
          [ 1.1884],                           [-0.5598],
          [ 0.2629]],                          [-0.6689]],

         [[-1.3980],                          [[-0.4329],
          [-0.5598],                           [ 1.1884],
          [-0.6689]],                          [ 0.2629]],

         [[-0.0211],                          [[-0.0211],
          [-0.5385],                           [-0.5385],
          [-0.0628]]]])                        [-0.0628]]]])
# torch.Size([1, 4, 3, 1])
'''
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

• 应用：可以用来分割图像块
  例如：将H×W×C的图片切分成N个P×P×C的图像块，其中序列长度N=H×W/P²，然后将每个图片转换成一维的向量表示P²C

x = rearrange(imge, 'b c (h h1) (w w2) -> b (h w) (h1 w2 c)', h1=p, w2=p)
1

更详细的功能可参照博客：https://blog.csdn.net/csdn_yi_e/article/details/109143580

有问题欢迎指正!(^o)/~