einops 0.4 : tutorial part 1 : 基本 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/25/2022 (0.4.1)

* 本ページは、einops の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Einops tutorial, part 1: basics

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

 

einops 0.4 : tutorial part 1 : 基本

Welcome to einops-land!

私達は次のようには書きません :

y = x.transpose(0, 2, 3, 1)

分かりやすいコードを書きます :

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

einops は (numpy, pytorch, chainer, gluon, tensorflow のような) 広く利用されているテンソル・パッケージをサポートし、それらを拡張します。

 

このチュートリアルの内容

• 基礎 : reordering, composition と decomposition of axes
• 演算 : rearrange, reduce, repeat
• 単一演算でどれほどできるのでしょう！

 

準備

# Examples are given for numpy. This code also setups ipython/jupyter
# so that numpy arrays in the output are displayed as images
import numpy
from utils import display_np_arrays_as_images
display_np_arrays_as_images()

 

遊ぶために画像のバッチをロードする

ims = numpy.load('./resources/test_images.npy', allow_pickle=False)
# There are 6 images of shape 96x96 with 3 color channels packed into tensor
print(ims.shape, ims.dtype)

(6, 96, 96, 3) float64

# display the first image (whole 4d tensor can't be rendered)
ims[0]

# second image in a batch
ims[1]

# 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(ims[0], 'h w c -> w h c')

 

軸の Composition

transposition は非常に一般的で有用ですが、einops により提供される他の機能に移りましょう :

# einops は batch と height を新しい height 次元にシームレスに compose することを可能にします。
# We just rendered all images by collapsing to 3d tensor!
rearrange(ims, 'b h w c -> (b h) w c')

# or batch と width で新しい次元を compose します。
rearrange(ims, 'b h w c -> h (b w) c')

# resulting dimensions are computed very simply
# length of newly composed axis is a product of components
# [6, 96, 96, 3] -> [96, (6 * 96), 3]
rearrange(ims, 'b h w c -> h (b w) c').shape

(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
rearrange(ims, 'b h w c -> (b h w c)').shape

(165888,)

 

軸の Decomposition

# 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

(2, 3, 96, 96, 3)

# finally, combine composition and decomposition:
rearrange(ims, '(b1 b2) h w c -> (b1 h) (b2 w) c ', b1=2)

# slightly different composition: b1 is merged with width, b2 with height
# ... so letters are ordered by w then by h
rearrange(ims, '(b1 b2) h w c -> (b2 h) (b1 w) c ', b1=2)

# 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)?
rearrange(ims, 'b h (w w2) c -> (h w2) (b w) c', w2=2)

 

軸の順序は重要

# compare with the next example
rearrange(ims, 'b h w c -> h (b w) c')

(96, 576, 3)

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

(96, 576, 3)

# what if b1 and b2 are reordered before composing to width?
rearrange(ims, '(b1 b2) h w c -> h (b1 b2 w) c ', b1=2) # produces 'einops'
rearrange(ims, '(b1 b2) h w c -> h (b2 b1 w) c ', b1=2) # produces 'eoipns'

 

einops.reduce に出会う

einops-land では何が起きたかを推測する必要はありません

x.mean(-1)

演算が何を行なうかを貴方が書くからです :

reduce(x, 'b h w c -> b h w', 'mean')

軸が出力に存在しない場合 — それを推測したように — 軸は reduce されました。

 

# average over batch
reduce(ims, 'b h w c -> h w c', 'mean')

(96, 96, 3)

# the previous is identical to familiar:
ims.mean(axis=0)
# but is so much more readable

(96, 96, 3)

# Example of reducing of several axes 
# besides mean, there are also min, max, sum, prod
reduce(ims, 'b h w c -> h w', 'min')

(96, 96)

# this is mean-pooling with 2x2 kernel
# image is split into 2x2 patches, each patch is averaged
reduce(ims, 'b (h h2) (w w2) c -> h (b w) c', 'mean', h2=2, w2=2)

(48, 288, 3)

# max-pooling is similar
# result is not as smooth as for mean-pooling
reduce(ims, 'b (h h2) (w w2) c -> h (b w) c', 'max', h2=2, w2=2)

(48, 288, 3)

# yet another example. Can you compute result shape?
reduce(ims, '(b1 b2) h w c -> (b2 h) (b1 w)', 'mean', b1=2)

(288, 192)

 

スタックと連結

# rearrange can also take care of lists of arrays with the same shape
x = list(ims)
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)
(6, 96, 96, 3)

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

(96, 96, 3, 6)

# that's equivalent to numpy stacking, but written more explicitly
numpy.array_equal(rearrange(x, 'b h w c -> h w c b'), numpy.stack(x, axis=3))

True

# ... or we can concatenate along axes
rearrange(x, 'b h w c -> h (b w) c').shape

(96, 576, 3)

# which is equivalent to concatenation
numpy.array_equal(rearrange(x, 'b h w c -> h (b w) c'), numpy.concatenate(x, axis=1))

True

 

軸の追加と削除

長さ 1 の新しい軸を作成するために 1 を書くことができます。同様にそのような軸を削除することもできます。

利用可能な同義語 () もあります。それはゼロ軸の composition でそれはまた unit length を持ちます。

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

(96, 576, 3)

 

要素を反復する

紹介する 3 番目の演算は repeat です。

# repeat along a new axis. New axis can be placed anywhere
repeat(ims[0], 'h w c -> h new_axis w c', new_axis=5).shape

(96, 5, 96, 3)

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

(96, 5, 96, 3)

# repeat along w (existing axis)
repeat(ims[0], 'h w c -> h (repeat w) c', repeat=3)

(96, 288, 3)

# repeat along two existing axes
repeat(ims[0], 'h w c -> (2 h) (2 w) c')

(192, 192, 3)

# order of axes matters as usual - you can repeat each element (pixel) 3 times 
# by changing order in parenthesis
repeat(ims[0], 'h w c -> h (w repeat) c', repeat=3)

(96, 288, 3)

Note repeat 演算は numpy.repeat, numpy.tile と同一の機能をカバーしますが実際にはそれ以上です。

 

Reduce ⇆ repeat

reduce と repeat は互いの反対のようなものです : 最初のものは要素の数を削減して、2 番目のものは増やします。

以下の例では各画像は最初に repeat され、元のテンソルに戻すために新しい軸に対して reduce します。演算パターンは互いの “reverse” であることに気付くでしょう。

repeated = repeat(ims, 'b h w c -> b h new_axis w c', new_axis=2)
reduced = reduce(repeated, 'b h new_axis w c -> b h w c', 'min')
assert numpy.array_equal(ims, reduced)

(6, 96, 2, 96, 3)
(6, 96, 96, 3)

 

Fancy なサンプル (順不同)

# interweaving pixels of different pictures
# all letters are observable
rearrange(ims, '(b1 b2) h w c -> (h b1) (w b2) c ', b1=2)

(192, 288, 3)

# interweaving along vertical for couples of images
rearrange(ims, '(b1 b2) h w c -> (h b1) (b2 w) c', b1=2)

(192, 288, 3)

# interweaving lines for couples of images
# exercise: achieve the same result without einops in your favourite framework
reduce(ims, '(b1 b2) h w c -> h (b2 w) c', 'max', b1=2)

(96, 288, 3)

# color can be also composed into dimension
# ... while image is downsampled
reduce(ims, 'b (h 2) (w 2) c -> (c h) (b w)', 'mean')

(144, 288)

# disproportionate resize
reduce(ims, 'b (h 4) (w 3) c -> (h) (b w)', 'mean')

(24, 192)

# spilt each image in two halves, compute mean of the two
reduce(ims, 'b (h1 h2) w c -> h2 (b w)', 'mean', h1=2)

(48, 576)

# split in small patches and transpose each patch
rearrange(ims, 'b (h1 h2) (w1 w2) c -> (h1 w2) (b w1 h2) c', h2=8, w2=8)

(96, 576, 3)

# stop me someone!
rearrange(ims, 'b (h1 h2 h3) (w1 w2 w3) c -> (h1 w2 h3) (b w1 h2 w3) c', h2=2, w2=2, w3=2, h3=2)

(96, 576, 3)

rearrange(ims, '(b1 b2) (h1 h2) (w1 w2) c -> (h1 b1 h2) (w1 b2 w2) c', h1=3, w1=3, b2=3)

(192, 288, 3)

# patterns can be arbitrarily complicated
reduce(ims, '(b1 b2) (h1 h2 h3) (w1 w2 w3) c -> (h1 w1 h3) (b1 w2 h2 w3 b2) c', 'mean', 
       h2=2, w1=2, w3=2, h3=2, b2=2)

(96, 576, 3)

# subtract background in each image individually and normalize
# pay attention to () - this is composition of 0 axis, a dummy axis with 1 element.
im2 = reduce(ims, 'b h w c -> b () () c', 'max') - ims
im2 /= reduce(im2, 'b h w c -> b () () c', 'max')
rearrange(im2, 'b h w c -> h (b w) c')

(96, 576, 3)

# pixelate: first downscale by averaging, then upscale back using the same pattern
averaged = reduce(ims, 'b (h h2) (w w2) c -> b h w c', 'mean', h2=6, w2=8)
repeat(averaged, 'b h w c -> (h h2) (b w w2) c', h2=6, w2=8)

(96, 576, 3)

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

(96, 576, 3)

# let's bring color dimension as part of horizontal axis
# at the same time horizontal axis is downsampled by 2x
reduce(ims, 'b (h h2) (w w2) c -> (h w2) (b w c)', 'mean', h2=3, w2=3)

(96, 576)

 

Ok, numpy is fun, しかし他のフレームワークではどのように einops を使用するのでしょう？

ims が numpy 配列であるとして遂行したものであるならば :

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

それが他のフレームワークに対してコードを適応させる方法です :

# pytorch:
rearrange(ims, 'b h w c -> w (b h) c')
# tensorflow:
rearrange(ims, 'b h w c -> w (b h) c')
# chainer:
rearrange(ims, 'b h w c -> w (b h) c')
# gluon:
rearrange(ims, 'b h w c -> w (b h) c')
# cupy:
rearrange(ims, 'b h w c -> w (b h) c')
# jax:
rearrange(ims, 'b h w c -> w (b h) c')

…well, you got the idea.

einops は総ての演算がフレームワークに native であるかのように逆伝播を可能にします。別のフレームワークに移行するとき演算は変わりません – einops 記法は普遍的です。

以上