Active Hackathon

Reinventing Deep Learning Operation Via Einops


Einops, an abbreviation of Einstein-Inspired Notation for operations is an open-source python framework for writing deep learning code in a new and better way. Einops provides us with new notation & new operations. It is  a flexible and powerful tool to ensure code readability and reliability with minimalist yet powerful API.

Supported Frameworks 


Sign up for your weekly dose of what's up in emerging technology.

Requirements & Installation


You can install Einops via PyPI.

!pip install einops


Here are a few examples to get started with Einops.

A. Basics of Einops

Instead of writing, y = x.transpose(0, 2, 3, 1), in Einops we represent it as y = rearrange(x, 'b c h w -> b h w c')

This demo will cover basics of Einops reordering, composition and decomposition of axes, operations like rearrange, reduce, repeat and what all you can do with these single operations.

Let’s get started!

  1. Import all the required packages and load the data. Here we are using 6 images. A sample of it is shown below:
  2. Perform all three operations:
  • Rearrange

Rearranging a single image:

 # 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 elements according to the pattern
 rearrange(ims[0], 'h w c -> w h c')

  Composition of axes: 

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

Decomposition of axes :

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

Order of the axes matter : 

 # compare with the next example
 rearrange(ims, 'b h w c -> h (b w) c')
 # order of axes in composition is different
 # rule is just as for digits in the number: leftmost digit is the most significant, 
 # while neighboring numbers differ in the rightmost axis.
 # you can also think of this as lexicographic sort
 rearrange(ims, 'b h w c -> h (w b) c') 

rearrange doesn’t change the number of elements and covers different numpy functions (like transpose, reshape, stack, concatenate, squeeze and expand_dims). You can check all types of rearrangement, here.

  • Reduce : Instead of worrying about  x.mean(-1),  Einops gives you a option of directly reducing the image as :
 # average over batch
 reduce(ims, 'b h w c -> h w c', 'mean')

If the axis is not present in the output, that means it is reduced and also provides different kinds of methods to reduce on like mean, min, max, sum, etc.

 # the previous is identical to familiar:
 # but is so much more readable 

reduce combines the same reordering syntax with reductions (mean, min, max, sum, prod, and any others). You can check all the utilities of the reduce function here.

  • Reduce ⇄ Repeat : reduce and repeat are like opposite of each other: first one reduces amount of elements, second one increases.
 # 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 additionally covers repeating and tiling. Some fancy examples are available here.

You can check Fundamental Demo of Einops in this Colab Notebook.

B. Einops for Deep Learning

This demo contains usage of some deep learning packages, important cases for DL models and at last functionality of einsops.asnumpy and einops.layers. You can select your choice of framework.Einops functions work with any tensor like they are native to the framework. For the example purpose, here the framework choice is PyTorch.

  1. Let’s start with very basic computation.

converting bchw to bhwc format and back:

 y = rearrange(x, 'b c h w -> b h w c')
  1. Backpropagation is quite common in DL models. Here is the code for backpropagation via einops.
 y0 = x
 y1 = reduce(y0, 'b c h w -> b c', 'max')
 y2 = rearrange(y1, 'b c -> c b')
 y3 = reduce(y2, 'c b -> ', 'sum')
 if flavour == 'tensorflow':
     print(reduce(tape.gradient(y3, x), 'b c h w -> ', 'sum'))
     print(reduce(x.grad, 'b c h w -> ', 'sum')) 
  1. einops.asnumpy : This function is used to convert tensor into numpy.
 from einops import asnumpy
 y3_numpy = asnumpy(y3)
 print(type(y3_numpy), y3_numpy.shape) 
  1. Common Building Blocks of Deep Learning
  • Flattening
 y = rearrange(x, 'b c h w -> b (c h w)')
  • Space-to-Depth 
 y = rearrange(x, 'b c (h h1) (w w1) -> b (h1 w1 c) h w', h1=2, w1=2)
  • Depth-to-Space
 y = reduce(x, 'b c h w -> b c', reduction='mean')
  1. Reductions
  • Global Average Pooling
 y = reduce(x, 'b c h w -> b c', reduction='mean')
  • Max-pooling, here with kernel 2 X 2
 y = reduce(x, 'b c (h h1) (w w1) -> b c h w', reduction='max', h1=2, w1=2)
         # you can skip names for reduced axes
 y = reduce(x, 'b c (h 2) (w 2) -> b c h w', reduction='max')
  1. Squeeze & 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') 
  1. Stacking

Start with a list of tensors

list_of_tensors = list(x)

New axis (one that enumerates tensors) appears first on the left side of expression. Just as if you were indexing a list – first you’d get tensor by index.

 tensors = rearrange(list_of_tensors, 'b c h w -> b h w c')

Or, Stacking along last dimension

 # or maybe stack along last dimension?
 tensors = rearrange(list_of_tensors, 'b c h w -> h w c b')
  1. Concatenation

    Concatenation over First Dimension

 tensors = rearrange(list_of_tensors, 'b c h w -> (b h) w c')

Concatenation over Last Dimension

 tensors = rearrange(list_of_tensors, 'b c h w -> h w (b c)')

You can check all other Functionalities here.

C. Einops for PyTorch


You can refer to official Codes, Docs & Tutorials here.

More Great AIM Stories

Aishwarya Verma
A data science enthusiast and a post-graduate in Big Data Analytics. Creative and organized with an analytical bent of mind.

Our Upcoming Events

Conference, in-person (Bangalore)
Machine Learning Developers Summit (MLDS) 2023
19-20th Jan, 2023

Conference, in-person (Bangalore)
Data Engineering Summit (DES) 2023
21st Apr, 2023

Conference, in-person (Bangalore)
MachineCon 2023
23rd Jun, 2023

3 Ways to Join our Community

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Telegram Channel

Discover special offers, top stories, upcoming events, and more.

Subscribe to our newsletter

Get the latest updates from AIM