einops rearrange说明

rearrang用于重新排列Tensor的维度

引用官网例子

# suppose we have a set of 32 images in "h w c" format (height-width-channel)
>>> images = [np.random.randn(30, 40, 3) for _ in range(32)]

image的shape为[32, 30, 40, 3]

# concatenate images along height (vertical axis), 960 = 32 * 30
>>> rearrange(images, 'b h w c -> (b h) w c').shape
(960, 40, 3)

维度的字母可以用任何字母表示，但是中间要用空格隔开，一一对应，这里b,h,w,c分别对应了32, 30, 40, 3
然后把它变换成(b h) w c, 括号里表示b和h这两个数字相乘，维度就变成了[bxh, w, c]

# concatenated images along horizontal axis, 1280 = 32 * 40
>>> rearrange(images, 'b h w c -> h (b w) c').shape
(30, 1280, 3)

把第2维h=30换到第一维，然后第2维换成bxw=32x40

# split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
>>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
(128, 15, 20, 3)

这里指定了h1=2, w1=2, 而(h1 h)=h1xh对应了30，所以算出h=15，同理根据w1xw = 40算出w=20，
输出维度为[bxh1xw1 h w c] = [32x2x2 15 20 3]