einops 이미지 패치 자르기

텐서 차원조작을 쉽게 해주는 패키지, einops

transformer 공부하며 코드 구현시 필수인게 이미지 패치를 자르는 것인데 이 패키지를 이용하면 매우 쉽게 이용할 수 있다.

torch나 tensorflow의 layer처럼 사용할 수도 있다고 한다.

 

import matplotlib.pyplot as plt
from PIL import Image

from einops import rearrange, reduce, repeat
from einops.layers.torch import Rearrange, Reduce # torch layer 처럼 사용가능
from einops.layers.tensorflow import Rearrange, Reduce

img = Image.open('./cat.jpg')

plt.imshow(img)

transform = Compose([Resize((224,224)), ToTensor()])
x = transform(img)
x = x.unsqueeze(0) # batch
x.shape

# 출력 
# torch.Size([1, 3, 224, 224])

patch_size = 16
patches = rearrange(x,
                    'b c (h s1) (w s2) -> b (h w) s1 s2 c',
                    s1=patch_size,
                    s2=patch_size)
patches.shape

# 출력
# torch.Size([1, 196, 16, 16, 3])

 

자른 이미지 패치 확인하기

plt.figure(figsize = (6,6))

for index in range(patches.shape[1]):
	# patches.shape[1] -> 196
    plt.subplot(14, 14, index+1)
    plt.imshow(patches[0][index])   # patch size = (B=1, 36, 16, 16,3 ) 
    # patches[0][index].shape -> torch.Size([16, 16, 3])
    plt.axis('off')

plt.show()

 

 

참고

이제 차원 관리는 einops