A new flavour of deep learning operations
Project description
einops
Flexible and powerful tensor operations for readable and reliable code. Supports numpy, pytorch, tensorflow, and others.
Contents
 Tutorial
 API microreference
 Installation
 Naming
 Why using einops
 Supported frameworks
 Contributing
 Github repository (for issues/questions)
Tutorial / Documentation
Tutorial is the most convenient way to see einops
in action (and right now works as a documentation)
 part 1: einops fundamentals
 part 2: einops for deep learning
 part 3: real code fragments improved with einops (so far only for pytorch)
Installation
Plain and simple:
pip install einops
einops
has no mandatory dependencies (code examples also require jupyter, pillow + backends).
To obtain the latest github version
pip install https://github.com/arogozhnikov/einops/archive/master.zip
API
einops
has minimalistic and powerful API.
Two operations provided (see einops tutorial for examples)
from einops import rearrange, reduce
# 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)
And two corresponding layers (einops
keeps separate version for each framework) with the same API.
from einops.layers.chainer import Rearrange, Reduce
from einops.layers.gluon import Rearrange, Reduce
from einops.layers.keras import Rearrange, Reduce
from einops.layers.torch import Rearrange, Reduce
from einops.layers.tensorflow import Rearrange, Reduce
Layers behave similarly to operations and have same parameters (for the exception of first argument, which is passed during call)
layer = Rearrange(pattern, **axes_lengths)
layer = Reduce(pattern, reduction, **axes_lengths)
# apply created layer to a tensor / variable
x = layer(x)
Example of using layers within a model:
# example given for pytorch, but code in other frameworks is almost identical
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Rearrange
model = Sequential(
Conv2d(3, 6, kernel_size=5),
MaxPool2d(kernel_size=2),
Conv2d(6, 16, kernel_size=5),
MaxPool2d(kernel_size=2),
# flattening
Rearrange('b c h w > b (c h w)'),
Linear(16*5*5, 120),
ReLU(),
Linear(120, 10),
)
Additionally two auxiliary functions provided
from einops import asnumpy, parse_shape
# einops.asnumpy converts tensors of imperative frameworks to numpy
numpy_tensor = asnumpy(input_tensor)
# einops.parse_shape gives a shape of axes of interest
parse_shape(input_tensor, 'batch _ h w') # e.g {'batch': 64, 'h': 128, 'w': 160}
Naming
einops
stays for EinsteinInspired Notation for operations
(though "Einstein operations" is more attractive and easier to remember).
Notation was loosely inspired by Einstein summation (in particular by numpy.einsum
operation).
Why using einops
notation
Semantic information (being verbose in expectations)
y = x.view(x.shape[0], 1)
y = rearrange(x, 'b c h w > b (c h w)')
while these two lines are doing the same job in some context,
second one provides information about input and output.
In other words, einops
focuses on interface: what is input and output, not how output is computed.
The next operation looks similar:
y = rearrange(x, 'time c h w > time (c h w)')
But it gives reader a hint: this is not an independent batch of images we are processing, but rather a sequence (video).
Semantic information makes code easier to read and maintain.
More checks
Reconsider the same example:
y = x.view(x.shape[0], 1) # x: (batch, 256, 19, 19)
y = rearrange(x, 'b c h w > b (c h w)')
second line checks that input has four dimensions, but you can also specify particular dimensions. That's opposed to just writing comments about shapes since comments don't work as we know
y = x.view(x.shape[0], 1) # x: (batch, 256, 19, 19)
y = rearrange(x, 'b c h w > b (c h w)', c=256, h=19, w=19)
Result is strictly determined
Below we have at least two ways to define depthtospace operation
# depthtospace
rearrange(x, 'b c (h h2) (w w2) > b (c h2 w2) h w', h2=2, w2=2)
rearrange(x, 'b c (h h2) (w w2) > b (h2 w2 c) h w', h2=2, w2=2)
there are at least four more ways to do it. Which one is used by the framework?
These details are ignored, since usually it makes no difference, but it can make a big difference (e.g. if you use grouped convolutions on the next stage), and you'd like to specify this in your code.
Uniformity
reduce(x, 'b c (x dx) > b c x', 'max', dx=2)
reduce(x, 'b c (x dx) (y dx) > b c x y', 'max', dx=2, dy=3)
reduce(x, 'b c (x dx) (y dx) (z dz)> b c x y z', 'max', dx=2, dy=3, dz=4)
These examples demonstrated that we don't use separate operations for 1d/2d/3d pooling, those all are defined in a uniform way.
Spacetodepth and depthto space are defined in many frameworks. But how about widthtoheight?
rearrange(x, 'b c h (w w2) > b c (h w2) w', w2=2)
Framework independent behavior
Even simple functions are defined differently by different frameworks
y = x.flatten() # or flatten(x)
Suppose x
shape was (3, 4, 5)
, then y
has shape ...
 numpy, cupy, chainer:
(60,)
 keras, tensorflow.layers, mxnet and gluon:
(3, 20)
 pytorch: no such function
Supported frameworks
Einops works with ...
 numpy
 pytorch
 tensorflow eager
 cupy
 chainer
 gluon
 tensorflow
 keras and tf.keras
 mxnet (experimental)
 jax (experimental)
Contributing
Best ways to contribute are
 share your feedback. Experimental APIs currently require thirdparty testing.
 spread the word about
einops
 if you like explaining things, alternative tutorials can be helpful
 translating examples in languages other than English is also a good idea
 finally, use
einops
notation in your papers to strictly define used operations!
Supported python versions
einops
works with python 3.5 or later.
There is nothing specific to python 3 in the code, we simply need to move further and the decision is not to support python 2.
Project details
Release history Release notifications  RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for einops0.2.0py2.py3noneany.whl
Algorithm  Hash digest  

SHA256  96b1bac57ddb591cccb927d24934d7601c3cdf3343a79a43d316a118d66e1043 

MD5  16af6b8b243d200a1c301c01d059713b 

BLAKE2b256  89325ded0a73d2e14ef5a6908a930c3e1e9f92ffead482a2f153182b7429066e 