Función elegante de transformación de dimensión _reorganizar

prefacio

Recientemente, cuando estaba mirando un código de proyecto de código abierto, a menudo vi que una función de este tipo se reorganizaba para realizar la conversión de dimensión en lugar de usar permutar. Aunque a veces permutar se puede reemplazar con reorganizar, no es tan legible como este último.

Vi esta publicación de blog cuando estaba escribiendo este artículo . ¿Qué hizo él para obtener la oferta de DeepMind? , Esta historia nos dice que la perseverancia en el aprendizaje y más rendimiento es eficaz.

Convertir canal

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])
Determinar si son iguales elemento por elemento: Verdadero

fusión de dimensiones

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)

antorcha.Tamaño([1, 9, 4])

Uso avanzado

De hecho, parece que no es lo mismo que el resultado de pixelshuffle, aunque las dimensiones son las mismas. Cuando tengas tiempo, vuelve a ir y venir...(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©: Verdadero