einops和einsum：直接操作张量的利器

einops和einsum：直接操作张量的利器

einops和einsum是Vision Transformer的代码实现里出现的两个操作tensor维度和指定tensor计算的神器，在卷积神经网络里不多见，本文将介绍简单介绍一下这两样工具，方便大家更好地理解Vision Transformer的代码。

einops：直接操作tensor维度的神器

github地址：https://github.com/arogozhnikov/einops

einops：灵活和强大的张量操作，可读性强和可靠性好的代码。支持numpy、pytorch、tensorflow等。

有了他，研究者们可以自如地操作张量的维度，使得研究者们能够简单便捷地实现并验证自己的想法，在Vision Transformer等需要频繁操作张量维度的代码实现里极其有用。

这里简单地介绍几个最常用的函数。

安装

einops的安装非常简单，直接pip即可：

pip install einops

rearrange

import torch
from einops import rearrange

i_tensor = torch.randn(16, 3, 224, 224)		# 在CV中很常见的四维tensor： （N，C，H，W）
print(i_tensor.shape)
o_tensor = rearrange(i_tensor, 'n c h w -> n h w c')
print(o_tensor.shape)

输出：

torch.Size([16, 3, 224, 224])
torch.Size([16, 224, 224, 3])

在CV中很常见的四维tensor：（N，C，H，W），即表示（批尺寸，通道数，图像高，图像宽），在Vision Transformer中，经常需要对tensor的维度进行变换操作，rearrange函数可以很方便地、很直观地操作tensor的各个维度。

除此之外，rearrange还有稍微进阶一点的玩法：

 
i_tensor = torch.randn(16, 3, 224, 224)
o_tensor = rearrange(i_tensor, 'n c h w -> n c (h w)')
print(o_tensor.shape)  
o_tensor = rearrange(i_tensor, 'n c (m1 p1) (m2 p2) -> n c m1 p1 m2 p2', p1=16, p2=16)
print(o_tensor.shape)  

输出：

torch.Size([16, 3, 50176])
torch.Size([16, 3, 14, 16, 14, 16])

可以进行指定维度的合并和拆分，注意拆分时需要在变换规则后面指定参数。

repeat

from einops import repeat

i_tensor = torch.randn(3, 224, 224)  
print(i_tensor.shape)
o_tensor = repeat(i_tensor, 'c h w -> n c h w', n=16)  
print(o_tensor.shape)

repeat时记得指定右侧repeat之后的维度值

输出：

torch.Size([3, 224, 224])
torch.Size([16, 3, 224, 224])

reduce

from einops import reduce

i_tensor = torch.randn((16, 3, 224, 224))
o_tensor = reduce(i_tensor, 'n c h w -> c h w', 'mean')
print(o_tensor.shape)
o_tensor_ = reduce(i_tensor, 'b c (m1 p1) (m2 p2)  -> b c m1 m2 ', 'mean', p1=16, p2=16)
print(o_tensor_.shape)

输出：

torch.Size([3, 224, 224])
torch.Size([16, 3, 14, 14])

reduce时记得指定左侧要被reduce的维度值

Rearrange

import torch
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Rearrange
 
model = Sequential(
    Conv2d(3, 64, kernel_size=3),
    MaxPool2d(kernel_size=2),
    Rearrange('b c h w -> b (c h w)'),      # 相当于 flatten 展平的作用
    Linear(64*15*15, 120), 
    ReLU(),
    Linear(120, 10)
)

i_tensor = torch.randn(16, 3, 32, 32)
o_tensor = model(i_tensor)
print(o_tensor.shape)

输出：

torch.Size([16, 10])

einops.layers.torch.Rearrange 是nn.Module的子类，可以放在网络里面直接当作一层。

torch.einsum：爱因斯坦简记法

爱因斯坦简记法：是一种由爱因斯坦提出的，对向量、矩阵、张量的求和运算 $\sum$ 的求和简记法。

在该简记法当中，省略掉的部分是：

1. 求和符号 $\sum$
2. 求和号的下标 $i$

省略规则为：默认成对出现的下标（如下例1中的 $i$ 和例2中的 $k$ ）为求和下标，被省略。

1）$x_i{y}_i$简化表示内积 $<\mathbf{x},\mathbf{y}>$
$$x_i{y}_i:=\sum _i{x}_i{y}_i=o$$

其中o为输出。

2. 用 $X_{ik}{Y}_{kj}$ 简化表示矩阵乘法 $\mathbf{X}\mathbf{Y}$
   $$X_{ik}{Y}_{kj}:=\sum _k{X}_{ik}{Y}_{kj}={\mathbf{O}}_{ij}$$
   其中 ${\mathbf{O}}_{ij}$ 为输出矩阵的第ij个元素。

这样的求和简记法，能够以一种统一的方式表示各种各样的张量运算（内积、外积、转置、点乘、矩阵的迹、其他自定义运算），为不同运算的实现提供了一个统一模型。

einsum在numpy和pytorch中都有实现，下面我们以在torch中为例，展示一下最简单的用法

import torch

i_a = torch.randn(16, 32, 4, 8)
i_b = torch.randn(16, 32, 8, 16)

out = torch.einsum('b h i j, b h j d -> b h i d', i_a, i_b)
print(out.shape)

输出：

torch.Size([16, 32, 4, 16])

可以看到，torch.einsum可以简便地指定tensor运算，输入的两个tensor维度分别为 $bhij$ 和 $bhjd$ ，经过tensor运算后，得到的张量维度为 $bhid$ 。代码运行结果与我们的预期一致。