einops的基本使用

一、einops的介绍

einops用于实现对张量（Tensor）的操作，它可以更简单的实现张量的变换，并且很容易上手。einops的作用类似pytorch中的review，transpose，permute等操作的合集。

二、安装与导包

pip install einops
from einops import rearrange,repeat,reduce

三、一些常用的操作

以下图为例，演示常用的操作。

3.1 rearrange操作

1. 维度交换

   from einops import rearrange
   # -> 前面的h w c代表img未变换前的shape，后面的w h c是要变成的shape
   # h w c只是人们的习惯,你完全可以写成 a b c
   rearrange(img,'h w c -> w h c')
   rearrange(img,'a b c -> b a c')
   

2. 增加/减少维度

   # 增加维度
   img7 = rearrange(img, 'h w c ->1 h w c')
   img7 = rearrange(img7, '1 h w c ->h w c')
   

3. 拆分单词（以w维度将单词拆分为6个字母）

   # 以w维度对图像进行分组
   # (b w)表示将第1维度的长度=b*w，b表示组数，w表示每组大小
   img2=rearrange(img,'h (b w) c->b h w c',b=6)
   

4. 将单词改为竖向

   # 先分组，而后将将b与h合并
   img3 = rearrange(img, 'h (b w) c->(b h) w c', b=6)
   

4. 先分组，再排成两行

# 将第1维度分成b1 b2 和w的乘积。两行可以确定b1=2
img4 = rearrange(img, 'h (b1 b2 w) c->(b1 h) (b2 w) c', b1=2,b2=3)

---

3.2 reduce操作

1. 求均值

# 求均值，->前后消失的符号代表求均值的维度
# 如果想要保留纬度数不变，可以用1带填补
img5=reduce(img.astype("float"), 'h w c -> h w', 'mean')
img5=reduce(img.astype("float"), 'h w c -> h w 1', 'mean')
img5=reduce(img.astype("float"), 'h w c -> h', 'mean')

2. 最小值

# 最小值
img6 = reduce(img.astype("float"), 'h w c -> h w 1', 'min')

其他计算：min, max, sum, prod

3.3 repeat操作

# 重复
img8=repeat(img, 'h w c -> (repeat h) w c', repeat=3)