目录
einops：易理解、易操作、跨框架的Tensor操作工具包
AI小甜心
2021年07月20日 16:25
收录于文集
共4篇

0.动机

张量（Tensor）的操作在不管是Tensorflow、Pytorch、还是Mxnet、Paddle中都是常用的操作。比如pytorch中的review，transpose，permute等操作。

而einops是提供常用张量（Tensor）操作的python包，提供了抽象于, , , , 

代码块
JavaScript
自动换行
复制代码
pytorch：y = x.transpose(0, 2, 3, 1)
einops：y = rearrange(x, 'b c h w -> b h w c')
复制成功

上述两端代码的效果一样：将x的size由（b,c,h,w）变为(b,h,w,c)。

1.安装：

代码块
JavaScript
自动换行
复制代码
pip install eniops
复制成功

2.一些常用的操作：

2.1 维度交换

代码块
JavaScript
自动换行
复制代码
from einops import rearrange
rearrange(ims[0], 'h w c -> w h c')
复制成功

2.2 以h维度压缩图片

代码块
JavaScript
自动换行
复制代码
rearrange(ims, 'b h w c -> (b h) w c')
复制成功

2.3 以w维度压缩图片

代码块
JavaScript
自动换行
复制代码
rearrange(ims, 'b h w c -> h (b w) c')
复制成功

2.4 重新定义batch的size：

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) h w c -> (b1 h) (b2 w) c ', b1=2)
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) h w c -> (b2 h) (b1 w) c ', b1=2)
复制成功

2.5 维度的顺序

代码块
JavaScript
自动换行
复制代码
rearrange(ims, 'b h w c -> h (b w) c')
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, 'b h w c -> h (w b) c')
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) h w c -> h (b1 b2 w) c ', b1=2) # produces 'einops'
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) h w c -> h (b2 b1 w) c ', b1=2)
复制成功

3.reduce操作：

代码块
JavaScript
自动换行
复制代码
reduce(ims, 'b h w c -> h w c', 'mean')
复制成功

代码块
JavaScript
自动换行
复制代码
#min, max, sum, prod
reduce(ims, 'b h w c -> h w', 'min')
复制成功

代码块
JavaScript
自动换行
复制代码
# 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)
复制成功

代码块
JavaScript
自动换行
复制代码
# 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)
复制成功

代码块
JavaScript
自动换行
复制代码
# yet another example. Can you compute result shape?
reduce(ims, '(b1 b2) h w c -> (b2 h) (b1 w)', 'mean', b1=2)
复制成功

4.stack和concatenate操作

代码块
JavaScript
自动换行
复制代码
# rearrange can also take care of lists of arrays with the same shape
x = list(ims)
rearrange(x, 'b h w c -> b h w c').shape
# ... or we can concatenate along axes
rearrange(x, 'b h w c -> h (b w) c').shape
复制成功

5.增加或者删除维度

代码块
JavaScript
自动换行
复制代码
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.repeat操作

代码块
JavaScript
自动换行
复制代码
repeat(ims[0], 'h w c -> h (repeat w) c', repeat=3)
复制成功

代码块
JavaScript
自动换行
复制代码
# 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)
复制成功

7.一些有意思的操作

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) h w c -> (h b1) (w b2) c ', b1=2)
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) h w c -> (h b1) (b2 w) c', b1=2)
复制成功

代码块
JavaScript
自动换行
复制代码
reduce(ims, '(b1 b2) h w c -> h (b2 w) c', 'max', b1=2)
复制成功

代码块
JavaScript
自动换行
复制代码
reduce(ims, 'b (h 2) (w 2) c -> (c h) (b w)', 'mean')
复制成功

代码块
JavaScript
自动换行
复制代码
reduce(ims, 'b (h 4) (w 3) c -> (h) (b w)', 'mean')
复制成功

代码块
JavaScript
自动换行
复制代码
reduce(ims, 'b (h1 h2) w c -> h2 (b w)', 'mean', h1=2)
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, 'b (h1 h2) (w1 w2) c -> (h1 w2) (b w1 h2) c', h2=8, w2=8)
复制成功

代码块
JavaScript
自动换行
复制代码
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)
复制成功

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

代码块
JavaScript
自动换行
复制代码
rearrange(ims, '(b1 b2) (h1 h2) (w1 w2) c -> (h1 b1 h2) (w1 b2 w2) c', h1=3, w1=3, b2=3)
复制成功

代码块
JavaScript
自动换行
复制代码
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)
复制成功

代码块
JavaScript
自动换行
复制代码
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')
复制成功

代码块
JavaScript
自动换行
复制代码
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)
复制成功

代码块
JavaScript
自动换行
复制代码
rearrange(ims, 'b h w c -> w (b h) c')
复制成功

代码块
JavaScript
自动换行
复制代码
reduce(ims, 'b (h h2) (w w2) c -> (h w2) (b w c)', 'mean', h2=3, w2=3)
复制成功

总结：

  • rearrange 仅仅只是改变Tensor的size，相关操比如：transpose, reshape, stack, concatenate, squeeze and expand_dims；

  • reduce操作的对象是张量的维度，或者维度的顺序，值，比如： mean, min, max, sum, prod；

  • repeat 执行的是复制之类的操作

  • composition and decomposition of axes are a corner stone, they can and should be used together

参考链接：

https://github.com/arogozhnikov/einops

python
本文为我原创，未经授权禁止转载
cv12239491