einops 튜토리얼

In [1]:

%config Completer.use_jedi = False
import numpy as np
from utils import display_np_arrays_as_images

• fundamentals: reordering, composition and decomposition of axes
• operations: rerrange, reduce, repeat

In [2]:

display_np_arrays_as_images()

In [3]:

img = np.load('test_images.npy', allow_pickle=False)
print(img.shape, img.dtype)

(6, 96, 96, 3) float64

In [4]:

img[0].shape

Out[4]:

(96, 96, 3)

In [5]:

img[0]

Out[5]:

In [6]:

img[1]

Out[6]:

In [7]:

from einops import rearrange, reduce, repeat

• rearrange: 요소를 재정렬
• height 와 width 바꿈
• 즉, 처음 두 축(차원)을 바꾼것

In [8]:

rearrange(img[0], 'h w c -> w h c')

Out[8]:

In [9]:

rearrange(img[0], 'h w c -> w h c').shape

Out[9]:

(96, 96, 3)

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!

In [10]:

rearrange(img, 'b h w c -> (b h) w c')

Out[10]:

• or compose a new dimension of batch and width

In [11]:

rearrange(img, 'b h w c -> h (b w) c')

Out[11]:

• resulting dimensions are computed very simply
• length of newly composed axis is a product of components
• [6, 96, 96, 3] -> [96, (6 * 96), 3]

In [12]:

print(img.shape)
print(rearrange(img, 'b h w c -> h (b w) c').shape)

(6, 96, 96, 3)
(96, 576, 3)

• we can compose more than two axes.
• let's flatten 4d array into 1d, resulting array has as many elements as the original

In [13]:

rearrange(img, 'b h w c -> (b h w c)').shape

Out[13]:

(165888,)

Decomposition of axis¶

• 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

In [14]:

rearrange(img, '(b1 b2) h w c -> b1 b2 h w c ', b1=2)

<array of shape (2, 3, 96, 96, 3)>

Out[14]:

• finally, combine composition and decomposition:

In [15]:

rearrange(img, '(b1 b2) h w c -> (b1 h) (b2 w) c ', b1=2)

Out[15]:

• slightly different composition: b1 is merged with width, b2 with height
• so letters are ordered by w then by h

In [16]:

rearrange(img, '(b1 b2) h w c -> (b2 h) (b1 w) c', b1=2)

Out[16]:

• 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)?

In [17]:

rearrange(img, 'b h (w w2) c -> (h w2) (b w) c', w2=2)

Out[17]:

Order of axes matters¶

In [18]:

rearrange(img, 'b h w c -> h (b w) c')

Out[18]:

• 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

In [19]:

rearrange(img, 'b h w c -> h (w b) c')

Out[19]:

• what if b1 and b2 are reordered before composing to width?

In [20]:

rearrange(img, '(b1 b2) h w c -> h (b1 b2 w) c', b1=2) # produces 'einops'

Out[20]:

In [21]:

rearrange(img, '(b1 b2) h w c -> h (b2 b1 w) c', b1=2) # produces 'eoipns'

Out[21]:

In [22]:

print(reduce(img, 'b h w c -> h w c', 'mean').shape)
reduce(img, 'b h w c -> h w c', 'mean')

(96, 96, 3)

Out[22]:

In [23]:

print(img.mean(axis=0).shape)
img.mean(axis=0)

(96, 96, 3)

Out[23]:

• Example of reducing of several axes
• besides mean, there are also min, max, sum, prod

In [24]:

reduce(img, 'b h w c -> h w', 'min')

Out[24]:

• this is mean-pooling with 2x2 kernel
• image is split into 2x2 patches, each patch is averaged

In [25]:

reduce(img, 'b (h h2) (w w2) c -> h (b w) c', 'mean', h2=2, w2=2)

Out[25]:

In [26]:

reduce(img, '(b1 b2) h w c -> (b2 h) (b1 w)', 'mean', b1=2)

Out[26]:

Stack and Concatenate¶

• rearrange can also take care of lists of arrays with the same shape

In [27]:

x = list(img)
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

<class 'list'> with 6 tensors of shape (96, 96, 3)

Out[27]:

(6, 96, 96, 3)

In [28]:

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

Out[28]:

(96, 96, 3, 6)

Addition or removal of axes¶

• You can write 1 to create a new axis of length 1. Similarly you can remove such axis.
• There is also a synonym () that you can use. That's a composition of zero axes and it also has a unit length.

In [29]:

x = rearrange(img, '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

(6, 1, 96, 96, 1, 3)
(6, 96, 96, 3)

• compute max in each image individuality, then show a difference

In [30]:

x = reduce(img, 'b h w c -> b () () c', 'max') - img
rearrange(x, 'b h w c -> h (b w) c')

Out[30]:

Repeating elements¶

• repeat along a new axis. New axis can be placed anywhere

In [31]:

repeat(img[0], 'h w c -> h new_axis w c', new_axis=5).shape

Out[31]:

(96, 5, 96, 3)

In [32]:

# shortcut
repeat(img[0], 'h w c -> h 5 w c').shape

Out[32]:

(96, 5, 96, 3)

• repeat along w (existing axis)

In [33]:

repeat(img[0], 'h w c -> h (repeat w) c', repeat=3)

Out[33]:

In [34]:

repeat(img[0], 'h w c -> h (3 w) c')

Out[34]:

• repeat along two existing axes

In [35]:

repeat(img[0], 'h w c -> (2 h) (2 w) c')

Out[35]:

• order of axes matters as usual - you can repeat each element (pixel) 3 times
• by changing order in parenthesis

In [36]:

repeat(img[0], 'h w c -> h (w 3) c')

Out[36]:

In [37]:

repeat(img[0], 'h w c -> h (3 w) c')

Out[37]:

In [ ]: