【详解】einops 优美的处理张量维度

einops：Deep learning operations reinvented (for pytorch, tensorflow, jax and others)

• 官方地址：http://einops.rocks/pytorch-examples.html
• einops works with python 3.6 or later.
• 以下代码所需数据请前往 此链接 免费下载

Einops works with …

• numpy
• pytorch
• tensorflow
• jax
• cupy
• chainer
• gluon
• tf.keras
• mxnet (experimental)

einops has a minimalistic yet powerful API. einops主要是 rearrange, reduce, repeat 这3个方法，(einops教程显示了覆盖堆叠、整形、换位、挤压/解压、重复、平铺、级联、视图和无数的缩减)

from einops import rearrange, reduce, repeat

# rearrange elements according to the pattern
output_tensor = rearrange(input_tensor, 't b c -> b c t')
# combine rearrangement and reduction
output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)
# copy along a new axis 
output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)

einops关注的是接口：关注输入和输出是什么，而不是如何计算输出（见下方代码）。第三行给读者一个提示：这不是我们正在处理的独立的一批图像，而是一个序列(视频)。einops 使得代码更易于阅读和维护。

y = x.view(x.shape[0], -1)
y = rearrange(x, 'b c h w -> b (c h w)')
y = rearrange(x, 'time c h w -> time (c h w)')

einops 基础

import numpy as np
from einops import rearrange, reduce, repeat

ims = np.load('test_images.npy', allow_pickle=False)
# There are 6 images of shape 96x96 with 3 color channels packed into tensor
print(ims.shape, ims.dtype)

(6, 96, 96, 3) float64

此 test_images.npy 可以类比成 一个 batch 的图像，batchsize=6，图像尺寸为 96x96，3 个channel。接下来的所有 einops操作都是对其进行处理。

einops主要是rearrange, reduce, repeat这3个方法，接下来将逐一解释。

1、rearrange：维度调整

rearrange(ims[0], 'h w c -> w h c')		# 调换维度

rearrange(ims, 'b h w c -> (b h) w c')  # 合并维度

# or compose a new dimension of batch and width
rearrange(ims, 'b h w c -> h (b w) c')
rearrange(ims, 'b h w c -> h (b w) c').shape

[6, 96, 96, 3] -> [96, (6 * 96), 3]

Decomposition of axis：

# decomposition is the inverse process - represent an axis as a combination of new axes
# several decompositions possible, so b1=2 is to decompose 6 to b1=2 and b2=3
rearrange(ims, '(b1 b2) h w c -> b1 b2 h w c ', b1=2).shape

(2, 3, 96, 96, 3)

此处的括号 (b1 b2) 将会按照 b1=2 b2自动计算，然后对原 b轴 进行拆解

2、reduce：

x.mean(-1)
reduce(x, 'b h w c -> b h w', 'mean')

# average over batch
reduce(ims, 'b h w c -> h w c', 'mean')



# this is mean-pooling with 2x2 kernel
# image is split into 2x2 patches, each patch is averaged
# 变小了
reduce(ims, 'b (h h2) (w w2) c -> h (b w) c', 'mean', h2=2, w2=2)

# yet another example. Can you compute result shape?
reduce(ims, '(b1 b2) h w c -> (b2 h) (b1 w)', 'mean', b1=2)

3、Addition or removal of axes：

x = rearrange(ims, 'b h w c -> b 1 h w 1 c') # functionality of numpy.expand_dims
print(x.shape)
print(rearrange(x, 'b 1 h w 1 c -> b h w c').shape) # functionality of numpy.squeeze
(6, 1, 96, 96, 1, 3) 
(6, 96, 96, 3)


# compute max in each image individually, then show a difference 
x = reduce(ims, 'b h w c -> b () () c', 'max') - ims
rearrange(x, 'b h w c -> h (b w) c')

4、Reductions

Simple global average pooling：

y = reduce(x, 'b c h w -> b c', reduction='mean')
y.shape
(torch.Tensor, torch.Size([10, 32]))

max-poolingwith a kernel 2x2：

y = reduce(x, 'b c (h h1) (w w1) -> b c h w', reduction='max', h1=2, w1=2)
y.shape
(torch.Tensor, torch.Size([10, 32, 50, 100]))

Squeeze and unsqueeze (expand_dims)：

# models typically work only with batches, 
# so to predict a single image ...
image = rearrange(x[0, :3], 'c h w -> h w c')
# ... create a dummy 1-element axis ...
y = rearrange(image, 'h w c -> () c h w')
# ... imagine you predicted this with a convolutional network for classification,
# we'll just flatten axes ...
predictions = rearrange(y, 'b c h w -> b (c h w)')
# ... finally, decompose (remove) dummy axis
predictions = rearrange(predictions, '() classes -> classes')

per-channel mean-normalization foreach image：

y = x - reduce(x, 'b c h w -> b c 1 1', 'mean')
y.shape
(torch.Tensor, torch.Size([10, 32, 100, 200]))

per-channel mean-normalization forwhole batch：

y = x - reduce(y, 'b c h w -> 1 c 1 1', 'mean')
y.shape
(torch.Tensor, torch.Size([10, 32, 100, 200]))

5、Concatenation：

concatenate over the first dimension：

tensors = rearrange(list_of_tensors, 'b c h w -> (b h) w c')
tensors.shape
(torch.Tensor, torch.Size([1000, 200, 32]))

6、Shuffling within a dimension：

channel shuffle：

y = rearrange(x, 'b (g1 g2 c) h w-> b (g2 g1 c) h w', g1=4, g2=4)
y.shape
(torch.Tensor, torch.Size([10, 32, 100, 200]))