einops用于实现对张量(Tensor)的操作,它可以更简单的实现张量的变换,并且很容易上手。einops的作用类似pytorch中的review,transpose,permute等操作的合集。
pip install einops
from einops import rearrange,repeat,reduce
以下图为例,演示常用的操作。
维度交换
from einops import rearrange
# -> 前面的h w c代表img未变换前的shape,后面的w h c是要变成的shape
# h w c只是人们的习惯,你完全可以写成 a b c
rearrange(img,'h w c -> w h c')
rearrange(img,'a b c -> b a c')
增加/减少维度
# 增加维度
img7 = rearrange(img, 'h w c ->1 h w c')
img7 = rearrange(img7, '1 h w c ->h w c')
拆分单词(以w维度将单词拆分为6个字母)
# 以w维度对图像进行分组
# (b w)表示将第1维度的长度=b*w,b表示组数,w表示每组大小
img2=rearrange(img,'h (b w) c->b h w c',b=6)
将单词改为竖向
# 先分组,而后将将b与h合并
img3 = rearrange(img, 'h (b w) c->(b h) w c', b=6)
# 将第1维度分成b1 b2 和w的乘积。两行可以确定b1=2
img4 = rearrange(img, 'h (b1 b2 w) c->(b1 h) (b2 w) c', b1=2,b2=3)
# 求均值,->前后消失的符号代表求均值的维度
# 如果想要保留纬度数不变,可以用1带填补
img5=reduce(img.astype("float"), 'h w c -> h w', 'mean')
img5=reduce(img.astype("float"), 'h w c -> h w 1', 'mean')
img5=reduce(img.astype("float"), 'h w c -> h', 'mean')
# 最小值
img6 = reduce(img.astype("float"), 'h w c -> h w 1', 'min')
其他计算:min, max, sum, prod
# 重复
img8=repeat(img, 'h w c -> (repeat h) w c', repeat=3)