Skip to main content

A new flavour of deep learning operations

Project description

einops package logo

einops

A new flavour of deep learning ops for numpy, pytorch, tensorflow, chainer, gluon, and others.

einops introduces a new way to manipulate tensors, providing safer, more readable and semantically richer code.



einops package logo

This video in better quality.

Documentation / Tutorials

Tutorials are the most convenient way to see einops in action

(Tutorials are working as a documentation too.)

Installation

einops has no mandatory dependencies.

To obtain the latest github version

pip install https://github.com/arogozhnikov/einops/archive/master.zip

pypi release will follow soon.

API

Micro-reference on public API.

einops API is very minimalistic and powerful.

Two operations provided (see the guide to einops fundamentals)

from einops import rearrange, reduce

# rearrange elements according to pattern
output_tensor = rearrange(input_tensor, pattern, **axes_lengths)

# rearrange elements according to pattern
output_tensor = reduce(input_tensor, pattern, reduction, **axes_lengths)

Two auxiliary functions

from einops import asnumpy, parse_shape

# einops.asnumpy converts tensors of imperative frameworks to numpy
numpy_tensor = asnumpy(input_tensor)

# einops.parse_shape returns a shape in the form of a dictionary, axis name mapped to its length 
parse_shape(input_tensor, pattern)

And two layers (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

Einops layers are behaving in the same way as operations, and have same parameters (for the exception of first argument, which should be passed during call)

layer = Rearrange(pattern, **axes_lengths)
# applying to tensor
x = layer(x)

layer = Reduce(pattern, reduction, **axes_lengths)
# applying to tensor
x = layer(x)

Usually it is more convenient to use layers, not operations, to build models

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

model = Sequential(
    Conv2d(3, 6, kernel_size=5),
    MaxPool2d(kernel_size=2),
    Conv2d(6, 16, kernel_size=5),
    Reduce('b c (h h2) (w w2) -> b (c h w)', 'max', h2=2, w2=2), # combined pooling and flattening
    Linear(16*5*5, 120), 
    ReLU(),
    Linear(120, 10), 
)

Layers are available for chainer, gluon, keras and torch.

Naming and terminology

einops stays for Einstein-Inspired Notation for operations (though "Einstein operations" sounds simpler and more attractive).

Notation was loosely inspired by Einstein summation (in particular by numpy.einsum operation).

  • Terms tensor and ndarray are equivalently used and refer to multidimensional array
  • Terms axis and dimension are also equivalent

Why using einops notation

Semantic information:

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 to previous two:

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

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

Back to 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 there are four dimensions in input, 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 may be understood differently within 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

  • spread the word about einops
  • prepare a guide/post/tutorial for your favorite deep learning framework
  • translating examples in languages other than English is also a good idea
  • use einops notation in your papers to strictly define an operation you're using

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 I decided 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.1.0.tar.gz (15.1 kB view details)

Uploaded Source

Built Distribution

einops-0.1.0-py3-none-any.whl (16.8 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: einops-0.1.0.tar.gz
  • Upload date:
  • Size: 15.1 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/1.11.0 pkginfo/1.4.2 requests/2.18.4 setuptools/39.1.0 requests-toolbelt/0.8.0 tqdm/4.23.4 CPython/3.6.6

File hashes

Hashes for einops-0.1.0.tar.gz
Algorithm Hash digest
SHA256 4fd64864fcb8159074da3213b9327c242536784416cbf423745ef8579850d30b
MD5 9ab814121e666c0cf50944dd1a892ee3
BLAKE2b-256 32d3efa7cd2f496ebe32b7b7e17bb0d91f2fb216b032ce572e63bab767f46e32

See more details on using hashes here.

File details

Details for the file einops-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: einops-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 16.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/1.11.0 pkginfo/1.4.2 requests/2.18.4 setuptools/39.1.0 requests-toolbelt/0.8.0 tqdm/4.23.4 CPython/3.6.6

File hashes

Hashes for einops-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 4ab512fe059c0841e1a315449ca9d7f35eaa05c8c095a14f2c1b92b2b77684d2
MD5 489ed693686d3b474d44e2959f0e9025
BLAKE2b-256 361fa5e6496a167b6e892123b201ccce251184457ce7a1b6c8e06662c09bfbab

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