python einops张量操作工具包

eniops是python提供的一个对张量维度进行随心所欲操作的包，具有十分强大的功能，下面就让我们来见识一下eniops的强大表达能力。

rearrange

具有交换张量维度的功能，下面的例子将C,H,W方式存储的图片转化为H,W,C的存储方式

import torch
from einops import rearrange

x=torch.rand((10,3,224,224))
print(x.size())

x = rearrange(x, 'b c h w ->b h w c') # 
print(x.size())

合并维度，下面的例子将H,W两个维度合并，即转化为一维向量

x=torch.rand((10,3,224,224))
print(x.size())

x = rearrange(x, 'b c h w ->b c (h w)') #
print(x.size())

合并batch,h,w维度

x=torch.rand((10,3,224,224))
print(x.size())

x = rearrange(x, 'b c h w ->c (b h w)') #
print(x.size())

 

对指定维度进行拆分，下面的例子将Batch维度进一步拆分为5x2

x = rearrange(x, '(b1 b2) c h w ->b1 b2 c h w',b2=2) #
print(x.size())

 Reduce

对指定维度求均值

from einops import reduce
x=torch.rand((10,3,224,224))
print(x.size())

x = reduce(x, 'b c h w -> b h w', reduction='mean')
print(x.size())

进行2x2平均池化

x=torch.rand((10,3,224,224))
print(x.size())

x =  reduce(x, 'b c (h h2) (w w2) -> b c h w', reduction='mean', h2=2, w2=2) # h 和 w 变小，实现2 x 2 的均值池化
print(x.size())

 全局最大池化

from einops import reduce
x=torch.rand((10,3,224,224))
print(x.size())

x =  reduce(x, 'b c h w -> b c ', reduction='max') # 
print(x.size())

 Repeat

在一个维度上进行拷贝

from einops import repeat
x=torch.rand((10,3,224,224))
print(x.size())

x =  repeat(x, 'b c h w -> r b c h w', r=100)# copy along a new axis
print(x.size())

其他

增加一个维度

x=torch.rand((10,3,224,224))
print(x.size())

x =  rearrange(x, 'b c h w -> 1 b c h w')#
print(x.size())

减少一个维度

x=torch.rand((10,1,3,224,224))
print(x.size())

x =  rearrange(x, 'b 1 c h w -> b c h w')#
print(x.size())

 

求最大值但保留维度

x=torch.rand((10,3,224,224))
print(x.size())

x =  reduce(x, 'b c h w -> b c () ()','max')#
print(x.size())

flatten展开

from einops import rearrange
from einops import reduce
from einops import repeat
x=torch.rand((10,3,224,224))
print(x.size())

x =  rearrange(x, 'b c h w -> (b c h w)')#
print(x.size())

patch 

将图片分为16x16的patch,共14x14=196个

from einops import rearrange
from einops import reduce
from einops import repeat
x=torch.rand((10,3,224,224))
print(x.size())

x =  rearrange(x, 'b c (h hp) (w wp) -> b c (h w) hp wp',hp=16,wp=16)#
print(x.size())

总结

• 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://zhuanlan.zhihu.com/p/372692913