einops.rearrange、repeat、reduce==＞对维度进行操作

1.einops.rearrange==》重新指定维度

def rearrange(tensor, pattern, **axes_lengths):

einops.rearrange is a reader-friendly smart element reordering for multidimensional tensors. This operation includes functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze, stack, concatenate and other operations.

import numpy as np
from einops import rearrange, repeat# 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)]
# stack along first (batch) axis, output is a single array ：(32, 30, 40, 3)
print(rearrange(images, 'b h w c -> b h w c').shape)# concatenate images along height (vertical axis), 960 = 32 * 30  ：(960, 40, 3)
print(rearrange(images, 'b h w c -> (b h) w c').shape)# concatenated images along horizontal axis, 1280 = 32 * 40  ：(30, 1280, 3)
print(rearrange(images, 'b h w c -> h (b w) c').shape)# reordered axes to "b c h w" format for deep learning  ：(32, 3, 30, 40)
print(rearrange(images, 'b h w c -> b c h w').shape)# flattened each image into a vector, 3600 = 30 * 40 * 3  ：(32, 3600)
print(rearrange(images, 'b h w c -> b (c h w)').shape)# ======================================================================================================================
# 这里(h h1) (w w1)就相当于h与w变为原来的1/h1,1/w1倍# split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2  ：(128, 15, 20, 3)
print(rearrange(images, 'b (h h1) (w w1) c -> (b h1 w1) h w c', h1=2, w1=2).shape)# space-to-depth operation  ：(32, 15, 20, 12)
print(rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape)

2.einops.repeat==》重排和重复（增加）维度

einops.repeat allows reordering elements and repeating them in arbitrary combinations. This operation includes functionality of repeat, tile, broadcast functions.

import numpy as np
from einops import rearrange, repeat,reduce# a grayscale image (of shape height x width)
image = np.random.randn(30, 40)# change it to RGB format by repeating in each channel：(30, 40, 3)
print(repeat(image, 'h w -> h w c', c=3).shape)# repeat image 2 times along height (vertical axis)：(60, 40)
print(repeat(image, 'h w -> (repeat h) w', repeat=2).shape)# repeat image 2 time along height and 3 times along width：(30, 120)
print(repeat(image, 'h w -> h (repeat w)', repeat=3).shape)# convert each pixel to a small square 2x2. Upsample image by 2x：(60, 80)
print(repeat(image, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape)# pixelate image first by downsampling by 2x, then upsampling：(30, 40)
downsampled = reduce(image, '(h h2) (w w2) -> h w', 'mean', h2=2, w2=2)
print(repeat(downsampled, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape)

3.einops.reduce

einops.reduce provides combination of reordering and reduction using reader-friendly notation.

import numpy as np
from einops import rearrange,reducex = np.random.randn(100, 32, 64)
# perform max-reduction on the first axis:(32, 64)
print(reduce(x, 't b c -> b c', 'max').shape) # same as previous, but with clearer axes meaning:(32, 64)
print(reduce(x, 'time batch channel -> batch channel', 'max').shape)x = np.random.randn(10, 20, 30, 40)
# 2d max-pooling with kernel size = 2 * 2 for image processing:(10, 20, 15, 20)
y1 = reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h2=2, w2=2)
print(y1.shape)# if one wants to go back to the original height and width, depth-to-space trick can be applied:(10, 5, 30, 40)
y2 = rearrange(y1, 'b (c h2 w2) h1 w1 -> b c (h1 h2) (w1 w2)', h2=2, w2=2)
print(y2.shape)# Adaptive 2d max-pooling to 3 * 4 grid:(10, 20, 3, 4)
print(reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h1=3, w1=4).shape)# Global average pooling:(10, 20)
print(reduce(x, 'b c h w -> b c', 'mean').shape)