einops使用方法

einops(爱因斯坦标记法)用法操作tensor。支持

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

官方文档: einops

einops让代码可读性更强，例如:

传统代码:

y = x.transpose(0, 2, 3, 1)

einops代码:

y = rearrange(x, 'b c h w -> b h w c')

相关代码:

import numpy
from utils import display_np_arrays_as_images
display_np_arrays_as_images()

ims = numpy.load('./resources/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

其中包含了5幅图像:


等等。

三种操作符:

# we'll use three operations
from einops import rearrange, reduce, repeat

# rearrange, as its name suggests, rearranges elements
# below we swapped height and width.
# In other words, transposed first two axes (dimensions)
rearrange(ims[0], 'h w c -> w h c')

Composition

# einops allows seamlessly composing batch and height to a new height dimension
# We just rendered all images by collapsing to 3d tensor!
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')

# 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

# finally, combine composition and decomposition:
rearrange(ims, '(b1 b2) h w c -> (b1 h) (b2 w) c ', b1=2)

# slightly different composition: b1 is merged with width, b2 with height
# ... so letters are ordered by w then by h
rearrange(ims, '(b1 b2) h w c -> (b2 h) (b1 w) c ', b1=2)

# move part of width dimension to height. 
# we should call this width-to-height as image width shrunk by 2 and height doubled. 
# but all pixels are the same!
# Can you write reverse operation (height-to-width)?
rearrange(ims, 'b h (w w2) c -> (h w2) (b w) c', w2=2)

比较传统的x.mean(-1), einops的操作更加直观:

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

if axis is not present in the output — you guessed it — axis was reduced.

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

与ims.mean(axis=0一样。

# Example of reducing of several axes 
# besides mean, there are also min, max, sum, prod
reduce(ims, 'b h w c -> h w', 'min')

# 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)

# max-pooling is similar
# result is not as smooth as for mean-pooling
reduce(ims, 'b (h h2) (w w2) c -> h (b w) c', 'max', 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)

# rearrange can also take care of lists of arrays with the same shape
x = list(ims)
print(type(x), 'with', len(x), 'tensors of shape', x[0].shape)
# that's how we can stack inputs
# "list axis" becomes first ("b" in this case), and we left it there
rearrange(x, 'b h w c -> b h w c').shape

'''
输出
 with 6 tensors of shape (96, 96, 3)
(6, 96, 96, 3)
'''

# but new axis can appear in the other place:
rearrange(x, 'b h w c -> h w c b').shape

# that's equivalent to numpy stacking, but written more explicitly
numpy.array_equal(rearrange(x, 'b h w c -> h w c b'), numpy.stack(x, axis=3))

# ... or we can concatenate along axes
rearrange(x, 'b h w c -> h (b w) c').shape

# which is equivalent to concatenation
numpy.array_equal(rearrange(x, 'b h w c -> h (b w) c'), numpy.concatenate(x, axis=1))

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

# 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')

# repeat along a new axis. New axis can be placed anywhere
repeat(ims[0], 'h w c -> h new_axis w c', new_axis=5).shape

# repeat along w (existing axis)
repeat(ims[0], 'h w c -> h (repeat w) c', repeat=3)

# repeat along two existing axes
repeat(ims[0], 'h w c -> (2 h) (2 w) c')

总结:

• rearrange doesn’t change number of elements and covers different numpy functions (like transpose, reshape, stack, concatenate, squeeze and expand_dims)
• reduce combines same reordering syntax with reductions (mean, min, max, sum, prod, and any others)
• repeat additionally covers repeating and tiling
• composition and decomposition of axes are a corner stone, they can and should be used together

# select one from 'chainer', 'gluon', 'tensorflow', 'pytorch' 
flavour = 'pytorch'

print('selected {} backend'.format(flavour))
if flavour == 'tensorflow':
    import tensorflow as tf
    tape = tf.GradientTape(persistent=True)
    tape.__enter__()
    x = tf.Variable(x) + 0
elif flavour == 'pytorch':
    import torch
    x = torch.from_numpy(x)
    x.requires_grad = True
elif flavour == 'chainer':
    import chainer
    x = chainer.Variable(x)
else:
    assert flavour == 'gluon'
    import mxnet as mx
    mx.autograd.set_recording(True)
    x = mx.nd.array(x, dtype=x.dtype)
    x.attach_grad()

转为numpy

from einops import asnumpy
y3_numpy = asnumpy(y3)

print(type(y3_numpy))

Sequeeze and unsequeeze

# 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')

list_of_tensors = list(x)

# Stacking over the first dimension?
tensors = rearrange(list_of_tensors, 'b c h w -> b h w c')

# stacking along last dimension?
tensors = rearrange(list_of_tensors, 'b c h w -> h w c b')

# Concatenation over the first dimension?
tensors = rearrange(list_of_tensors, 'b c h w -> (b h) w c')

# Split a dimension

bbox_x, bbox_y, bbox_w, bbox_h = 
    rearrange(x, 'b (coord bbox) h w -> coord b bbox h w', coord=4, bbox=8)
# now you can operate on individual variables
max_bbox_area = reduce(bbox_w * bbox_h, 'b bbox h w -> b h w', 'max')
guess(bbox_x.shape)
guess(max_bbox_area.shape)