1
import torch
2
from einops import rearrange,reduce,repeat
3

4
x = torch.randn(2,3,4,5)
5
# 1 transpose
6
out1 = x.transpose(1,2)
7
out2 = rearrange(x,'b i h w ->b h i w')
8

9
# 2 reshape
10
out1 = x.reshape(-1,4,5)
11
out2 = rearrange(x,'b i h w->(b i) h w')
12

13
out3 = rearrange(out2,'(b i) h w -> b i h w',b=2)  # b=2
14
flag = torch.allclose(out3,x)
15
print(flag)
16

17
# 池化操作
18
out1 = reduce(x,'b i h w -> b i h','mean') # avg pool
19
out2 = reduce(x,'b i h w -> b i h 1','sum')  # keep dimension
20
out3 = reduce(x,'b i h w-> b i','max')
21

22
# repate操作
23
out1 = rearrange(x,'b i h w -> b i h w 1') #extend dim  torch.unsqueeze
24
# print(out1)
25
out2 = repeat(out1,'b i h w 1 -> b i h w 2') # torch.tile
26

27
out3 = repeat(x,'b i h w -> b i (2 h) (2 w)')