einops的简单使用

简介

einops是一个用于操作张量的库，它的出现可以替代我们平时使用的reshape、view、transpose和permute等操作，相信接触过深度学习的同学一定对这些函数比较的熟悉。

einops相较于上面说的那些函数，最显著的区别就是逻辑更加的清晰，用网上的一句话来说的话，就是可以避免view、transpose等函数的神秘主义🤦‍♂️

第一次看见别人用einops这个库是在看ViT代码的时候，在简单的使用之后，感觉确实会比一般的维度操作工具更加的方便和直观，所以这里就einops的使用做一个简单的介绍

使用

einops操作的对象是张量，这里的张量可以是python中的矩阵，也可以是Pytorch和Tensorflow中的tensor，所以这里我就先读取一张图片

import cv2

path = './cat.jpg'
image = cv2.imread(path)
print(image.shape)
cv2.imshow('image', image)
cv2.waitKey(0)

运行上面的代码，图片就展示出来了😺

图片的大小：

接下来einops就可以对这张图片进行操作，其中einops常用的方法就三个，分别是rearrange, repeat, reduce

其中rearrange是对张量进行维度的变换，类似于之前的reshape和view；repeat主要用于张量的维度复制；reduce可以实现一些上采样和下采样的操作

我们这里先将这三个方法import进来

from einops import rearrange, repeat, reduce

rearrange

rearrange的函数定义如下：

def rearrange(tensor, pattern: str, **axes_lengths):

其中tensor就是我们想要操作的张量，这里就是我们那张猫的图片，最关键的是第二个参数pattern，这里列出我们操作这张图片的代码：

image = rearrange(image, 'h w c -> w h c')

可以看出pattern就是一串字符串，这个字符串里面就包含了我们的操作信息，h w c -> w h c这个信息很形象，原始的图片大小是(453. 512, 3)，分别对应了图片的高，宽和通道数，h w c -> w h c就是将高和宽互换，我们运行这行代码，就可以得到操作后的图像

可以看到操作成功了，我们还可以随意操作一下：

image = rearrange(image, 'h (a w) c -> w (a h) c', a=2)

值得注意的是，我们这里引入了一个参数a，在引入参数后，必须在后面定义一下参数的值，否则会报错。运行上面的代码后，就可以得到结果

这个其实也很好理解，原本的参数是h (a w) c，分别对应了高、宽和通道数，其中a=2，所以这里的w就是图片真实宽度的1/2，相当于宽被切一半，然后再重新排列成w (a h) c， 表示将切开了两块高和宽互换再拼接再一起，就得到了图片中的结果。同样的你可以将a设置为4，得到的就是下面的结果：

这样子的操作就会十分的清晰，而不是原先了类似transpose(0, 1)那样，不知道到底在操作哪个维度。

repeat

了解完一个函数的操作之后，后面两个其实大同小异，只是操作完的效果不一样。

repeat主要用在将张量在某一个维度进行复制，就和它的名字一样。我们执行以下的代码：

image = repeat(image, 'h w c -> h (repeat w) c', repeat=2)

得到的结果：

因为我们在宽度的维度后面加了个参数repeat，所以最终的结果就是在宽度这个维度重复，重复的次数取决于我们repeat的值，我们这里设置为了2

同样的我们可以在高度这个维度进行重复：

image = repeat(image, 'h w c -> (repeat h) w c', repeat=2)

reduce

reduce能实现的效果就有些多，这里也对这些操作一一的展示

• 求平均值mean

  利用如下的代码，就可以求我们想要维度的平均值：

  from einops import rearrange, repeat, reduce
  import cv2
  import numpy as np
  
  # read image
  path = './cat.jpg'
  image = cv2.imread(path).astype(np.float32) / 255.
  
  # einops
  image = reduce(image, 'h w c -> h w', reduction='mean')
  print(image.shape)
  
  # show image
  cv2.imshow('image', image)
  cv2.waitKey(0)
  

  其中操作张量的代码是

  image = reduce(image, 'h w c -> h w', reduction='mean')
  

  这里将图片转换为了float类型并归一化了，是因为reduce求平均值不能对整型的数据操作。

  我们在看回这行代码，发现前面的h w c到后面就变成了h w，这其实就代表了我们是在通道这个维度对图片进行了求平均值的操作，求完平均值之后，图片也就变成了二维的矩阵

  

  同理，可以对高和宽进行求平均值的操作，但是求完之后就不再是标准的图片，所以也就无法展示出来了。

  在深度学习里面，输入的数据往往还会多一个维度B，也就是我们平常说的Batch size，运用reduce的求平均值操作也可以很轻松的在B这个维度对一个批次的所有数据求平均值

• 另外一种mean

  我们利用reduce还可以实现将图片切分，然后堆叠在一起求平均值、最大最小值等操作

  这里我直接使用了numpy去生成一个矩阵，大小为(1, 3, 224, 224)

  代码如下：

  import numpy as np
  from einops import rearrange, repeat, reduce
  
  # image
  image = np.random.random((1, 3, 224, 224))
  
  # einops
  image = reduce(image, 'b c (h h2) (w w2) -> b c h w', reduction='mean', h2=4, w2=4)
  
  print(image.shape)
  

  代码运行出的结果：

  

  可以看到长宽高都缩小了4倍，更改reduction后面的操作就可以实现不同类型的pooling

总结

einops对张量的操作相较于之前的用0123来指定维度的方法会更加的直观，正如网上的大牛的评价一样：

einops正在缓慢而有力地渗入我代码的每一个角落和缝隙。如果你发现自己困扰于一堆高维的张量，这可能会改变你的生活。

后期如果了解了更详细的使用方法，再添加吧