einops rearrange 공부해보기

이번에 ViT를 공부하는 와중에 텐서를 조금 더 내 자유자재로 만들 수 있는 라이브러리를 발견했다

 

그건 바로 einops!

 

terminal 에 pip install einops    하면은 금방 설치할 수 있다.

 

한번 코드 예시를 통해 알아보겠다. 매우 직관적이므로 쉽게 알 수 있다.

 

from einops import rearrange
import torch.nn as nn
import torch

img = torch.rand(8, 224, 224, 3)

print(img.shape)
print(rearrange(img, 'b h w c -> b c h w').shape)
print(rearrange(img, 'b h w c -> (b c) h w').shape)
print(rearrange(img, 'b h (w1 w2) (c1 c2) -> b h c1 c2 w1 w2', w1 = 2, c1 = 1).shape)

 

torch.Size([8, 224, 224, 3])
torch.Size([8, 3, 224, 224])
torch.Size([24, 224, 224])
torch.Size([8, 224, 1, 3, 2, 112])

첫번째 결과값은 처음에 설정한 이미지의 크기이다.

두번째는 그 순서들을 살짝 변경했음을 알 수 있다.  이 부분은 reshape이나 view를 통해서도 할 수 있다.

세번째는 결과값에 대해 괄호로 묶어 곱셈으로 표현하였음을 알 수 있다.

네번째 결과값은 매우 신기하다. input 값들에 대해 괄호로 묶고 따로 인자를 주면은 텐서크기를 조절할 수 있다.

여기서 w1, c1 만 인자로 주어도 나머지 w2, c2같은 경우 알아서 배정이 된다.

 

앞으로 더 쓸일이 많을 거 같다.