Elegant dimension transformation _rearrange function

foreword

Recently, when I was looking at some open source project code, I often saw such a function rearrange to perform dimension conversion instead of using permute. Although sometimes permute can be replaced with rearrange, it is not as readable as the latter.

I saw this blog post when I was writing this article . What did he rely on to get DeepMind's offer? , This story tells us that perseverance in learning, and more output is effective.

Convert channel

import torch
from einops import rearrange 


# H, W, C
a = torch.randn(2, 2, 3)

# H, W, C -> C, H, W
a_permute = a.permute(2, 0, 1)
print('a_permute.shape:  ', a_permute.shape)

a_rearrange = rearrange(a, 'h w c -> c h w')
print('a_rearrange.shape:', a_rearrange.shape)

print('逐元素进行判断是否相等: ', a_permute.equal(a_rearrange))

a_permute.shape: torch.Size([3, 2, 2])
a_rearrange.shape: torch.Size([3, 2, 2])
Determine whether they are equal element by element: True

Dimension merge

import torch
from einops import rearrange


# B, C, H, W
a = torch.arange(9 * 2 * 2).view(1, 9, 2, 2)
# print(a)

b = rearrange(a, 'b c h w -> b c (h w)')
print(b.shape)

torch.Size([1, 9, 4])

Advanced usage

In fact, it seems that it is not the same as the pixelshuffle result, although the dimensions are the same. When you have time, go back and forth again...(todo)

import torch
from einops import rearrange


# B, C, H, W
a = torch.arange(36).view(1, 9, 2, 2)
# print(a)

# 建议在 torch1.12.x 测试 PixleShuffle这个类
# ps = torch.nn.PixleShuffle(3)

b = rearrange(a, 'b (c h1 w2) h w -> b c (h1 h) (w2 w)', h1=3, w2=3)
# print(b)
# b_ps = ps(a)
# print('b.equal(b_ps): ', b.equal(b_ps))

c = rearrange(b, 'b c (h1 h) (w2 w) -> b (c h1 w2) h w', h1=3, w2=3)
print('a.equal(c):', a.equal(c))

a.equal©: True