Skip to main content

A new flavour of deep learning operations

Project description

einops

Build Status PyPI version

Flexible and powerful tensor operations for readable and reliable code. Supports numpy, pytorch, tensorflow, and others.

Contents

Tutorial / Documentation

Tutorial is the most convenient way to see einops in action (and right now works as a documentation)

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 Einstein-Inspired 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 depth-to-space operation

# depth-to-space
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.

Space-to-depth and depth-to space are defined in many frameworks. But how about width-to-height?

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 ...

Contributing

Best ways to contribute are

  • share your feedback. Experimental APIs currently require third-party 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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

einops-0.2.0.tar.gz (15.6 kB view details)

Uploaded Source

Built Distribution

einops-0.2.0-py2.py3-none-any.whl (18.7 kB view details)

Uploaded Python 2Python 3

File details

Details for the file einops-0.2.0.tar.gz.

File metadata

  • Download URL: einops-0.2.0.tar.gz
  • Upload date:
  • Size: 15.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.1.1 pkginfo/1.5.0.1 requests/2.21.0 setuptools/42.0.2 requests-toolbelt/0.9.1 tqdm/4.37.0 CPython/3.6.5

File hashes

Hashes for einops-0.2.0.tar.gz
Algorithm Hash digest
SHA256 165ee28bcb60e5c2cbb801b5c78e181548ff8daa7c8fcabae5b251e55f7fe614
MD5 1528f3cea8a32f309abd5c89c4a351c3
BLAKE2b-256 d120df4006976ecb81f5e0bea3f5cbdbd90c4a70d5552b7541a82419449b9481

See more details on using hashes here.

File details

Details for the file einops-0.2.0-py2.py3-none-any.whl.

File metadata

  • Download URL: einops-0.2.0-py2.py3-none-any.whl
  • Upload date:
  • Size: 18.7 kB
  • Tags: Python 2, Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.1.1 pkginfo/1.5.0.1 requests/2.21.0 setuptools/42.0.2 requests-toolbelt/0.9.1 tqdm/4.37.0 CPython/3.6.5

File hashes

Hashes for einops-0.2.0-py2.py3-none-any.whl
Algorithm Hash digest
SHA256 96b1bac57ddb591cccb927d24934d7601c3cdf3343a79a43d316a118d66e1043
MD5 16af6b8b243d200a1c301c01d059713b
BLAKE2b-256 89325ded0a73d2e14ef5a6908a930c3e1e9f92ffead482a2f153182b7429066e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page