einops 张量操作

 

pip install einops

 

from einops import rearrange, reduce, repeat # 按给出的模式重组张量

output_tensor = rearrange(input_tensor, 't b c -> b c t') # 结合重组（rearrange）和reduction操作

output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2) # 沿着某一维复制

output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)

 

重新考虑和上面相同的例子:

y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
y = rearrange(x, 'b c h w -> b (c h w)')

•  

第二行检查了输入数据拥有四个维度（当然你也可以指定其他数字）
这和仅仅写注释标明数据维度是很不一样的，毕竟据我们所知，注释不能运行也无法阻止错误发生

y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
y = rearrange(x, 'b c h w -> b (c h w)', c